#!/usr/bin/env python3

import argparse
import logging
import os
import re
import requests
import simplediskimage

logging.basicConfig(level=logging.DEBUG)

def get_local_file(path, access):
    with open(path, access) as f:
        contents = f.read()
    return contents

def get_file(path):
    if re.search(r'https?://', path):
        request = requests.get(path, allow_redirects=True)
        request.raise_for_status()
        filename = path.split('/')[-1]
        #get_loca_file(path, 'wb')
        with open(path, 'wb') as f:
            contents = f.read()
    elif os.path.exists(path):
        #get_local_file(path, 'r')
        with open(path, 'r') as f:
            contents = f.read()
    else:
        raise Exception(f"Path {path} not found")
    print("satan: " + contents)
    return contents

def main(args):
    mlo = get_file(args.get('mlo', None))
    uboot = get_file(args.get('uboot', None))
    image = simplediskimage.DiskImage("foo.img", partition_table='msdos',
                                partitioner=simplediskimage.Sfdisk)
    pf = image.new_partition("fat16", partition_flags=["BOOT"])
    pf.copy(uboot)
    pf.copy(mlo)
    #pf.set_extra_bytes(45 * diskimage.SI.Mi)
    pf.set_fixed_size_bytes(48272 * simplediskimage.SI.ki)

    image.commit()
    print("sudo kpartx -av foo.img")
    print("...")
    print("sudo kpartx -dv foo.img")

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--mlo", required=True,
                        help="url or path to MLO file")
    parser.add_argument("--uboot", required=True,
                        help="url or path to u-boot.img file")
    args = vars(parser.parse_args())
    if args:
        main(args)
    else:
        exit(1)