#!/usr/bin/env python
# Copyright (C) 2019 Fabien Parent <fparent@baylibre.com>

import random
import zlib
import struct
import argparse
import sys

class UBootEnv:
    def __init__(self, env_size):
        self.env = [];
        self.env_size = env_size

    def __init__(self, env_size, env_file):
        self.env = [];
        self.env_size = env_size
        with open(env_file, "r") as env:
            self.env = env.readlines()

    def add(self, name, value):
        self.env.append("{}={}".format(name, value))

    def update(self, name, value):
        new_env = []
        for line in self.env:
            if line.startswith(name + '='):
                new_env.append("{}={}".format(name, value))
            else:
                new_env.append(line)
        self.env = new_env

    def write_binary(self, filename = "u-boot-env.bin"):
        with open(filename, "w+b") as out:
            out.seek(4)

            # Write environment
            for line in self.env:
                if line == '\n':
                    continue
                data = line.rstrip("\n").encode() + b'\0'
                out.write(data)

            while out.tell() < self.env_size:
                out.write(chr(0x00))

            # Compute CRC
            out.seek(4)
            crc = zlib.crc32(out.read(self.env_size - 4)) & 0xffffffff

            # Write CRC
            out.seek(0)
            out.write(struct.pack("I", crc))

    def gen_mac_addr(self, oui, num_iface):
        for i in range(num_iface):
            macaddr =  "{}:{:02X}:{:02X}:{:02X}".format(
                            oui,
                            random.randint(0, 255),
                            random.randint(0, 255),
                            random.randint(0, 255))
            varname = "eth{}addr".format(i)
            if i == 0:
                varname = "ethaddr"
            self.add(varname, macaddr)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input', default = "u-boot-initial-env",
                        help = 'Path to U-boot environment file')
    parser.add_argument('-s', '--size', default = 4096, type = int,
                        help = 'Size of the U-boot environment file')
    parser.add_argument('-o', '--output', default = "u-boot-env.bin",
                         help = 'Path to U-boot binary environment')
    parser.add_argument('--oui',
                        help = 'OUI to use when generating MAC addresses')
    parser.add_argument('--num-ethaddr', default = 0, type = int,
                        help = "Number of ethaddr to add to the environment")
    args = parser.parse_args()

    print("Initial environment: {}".format(args.input))
    print("Binary output environment: {}".format(args.output))
    print("OUI: {}".format(args.oui))
    print("Size: {}".format(args.size))

    if args.num_ethaddr > 0 and args.oui == None:
        sys.stderr.write("\nerror: --oui must be specified when --num-ethaddr is set\n")
        parser.print_help(sys.stderr)
        sys.exit(1)

    env = UBootEnv(args.size, args.input)
    if args.num_ethaddr > 0:
        env.gen_mac_addr(args.oui, args.num_ethaddr)
    env.write_binary(args.output)