#!/usr/bin/env python3
import struct, binascii, os

from Crypto.PublicKey import RSA

BOOTARGS_ENV = """bootdelay=0
verify=n
baudrate=115200
ipaddr=192.168.1.10
serverip=192.168.1.1
netmask=255.255.255.0
bootfile=uImage
phy_intf=mii,rgmii
use_mdio=0,1
phy_addr=2,3
gmac_debug=0
bootcmd=mmc read 0 0x7D400000 0x6E000 0x4000;bootm 0x7D400000;mmc read 0 0x1FFBFC0 0x50000 0xC800;fatload usb 0:1 0x05000000 initramfs_patched.uimg;bootm 0x02001FC0 0x05000000
bootargs=enforcing=0 androidboot.selinux=permissive androidboot.serialno=0123456789 androidboot.optimize=true console=ttyAMA0,115200 loglevel=8 blkdevparts=mmcblk0:1M(fastboot),1M(bootargs),20M(recovery),2M(deviceinfo),8M(baseparam),8M(pqparam),20M(logo),20M(logobak),40M(fastplay),40M(fastplaybak),40M(kernel),20M(misc),40M(trustedcore),1136M(system),400M(vendor),824M(cache),50M(private),8M(securestore),100M(customer),-(userdata) hbcomp=/dev/block/mmcblk0p14
bootargs_512M=mem=512M mmz=ddr,0,0,32M vmalloc=500M
bootargs_768M=mem=768M mmz=ddr,0,0,32M vmalloc=500M
bootargs_1G=mem=969M mmz=ddr,0,0,56M vmalloc=500M
bootargs_2G=mem=1961M mmz=ddr,0,0,56M vmalloc=500M
bootargs_3840M=mem=1961M mmz=ddr,0,0,56M vmalloc=500M
stdin=serial
stdout=serial
stderr=serial
"""

def env_block(size=65536):
    data = b""
    for line in BOOTARGS_ENV.strip().splitlines():
        line = line.strip()
        if line:
            data += line.encode('ascii') + b'\x00'
    data += b'\x00'
    raw = bytearray(b'\x00\x00\x00\x00' + data)
    raw += bytearray(size - len(raw))
    struct.pack_into('<I', raw, 0, binascii.crc32(bytes(raw[4:])) & 0xffffffff)
    return bytes(raw)

def main():
    os.makedirs('/root/pixboot', exist_ok=True)
    keyf = '/root/pixboot/rsa_key.pem'
    rsa = RSA.generate(2048, e=3)
    open(keyf, 'wb').write(rsa.exportKey('PEM'))
    print('generated new RSA-2048 e=3 key')
    pub = rsa.publickey()
    n = pub.n.to_bytes(256, 'big')
    e = struct.pack('<I', pub.e)
    crc = struct.pack('<I', binascii.crc32(n + e) & 0xffffffff)
    open('/root/pixboot/root_rsa_pub_crc.bin', 'wb').write(n + e + crc)
    open('/root/pixboot/bootargs.bin', 'wb').write(env_block())
    print('wrote root_rsa_pub_crc.bin (264B), bootargs.bin (64KB)')

if __name__ == '__main__':
    main()