#!/usr/bin/env python3

import os
import re
import shutil
import subprocess
##
import git  # https://pypi.org/project/GitPython/
import patch_ng  # https://pypi.org/project/patch-ng/


patches = ['git-version.patch', 'banner.patch', 'efi-iso.patch']

monkeypatch = {
    'config/general.h': {
        'enable': [
            'NET_PROTO_IPV6',
            'DOWNLOAD_PROTO_HTTPS',
            'DOWNLOAD_PROTO_FTP',
            'DOWNLOAD_PROTO_NFS',
            'HTTP_AUTH_NTLM',
            # 'HTTP_ENC_PEERDIST',
            # 'HTTP_HACK_GCE',
            'NSLOOKUP_CMD',
            'TIME_CMD',
            'DIGEST_CMD',
            'LOTEST_CMD',
            'VLAN_CMD',
            # 'PXE_CMD',  # Causes EFI to fail.
            'REBOOT_CMD',
            'POWEROFF_CMD',
            # 'IMAGE_NBI',  # I *think* this causes EFI to fail. Can I build it directly as a target w/o enabling?
            # 'IMAGE_ELF',  # Causes EFI to fail, and is auto-enabled for what needs it.
            # 'IMAGE_MULTIBOOT',  # Also enabled by default for MBR builds
            # 'IMAGE_SCRIPT',  # Enabled where needed etc.
            # 'IMAGE_BZIMAGE',  # http://lists.ipxe.org/pipermail/ipxe-devel/2017-March/005510.html
            # 'IMAGE_COMBOOT',  # Not really necessary since iPXE has native menus.
            # 'IMAGE_EFI',  # Enabled by default for EFI builds, and causes build error on bin/ipxe.dsk.tmp
            # 'IMAGE_PXE',  # EFI builds fail with this. related to PXE_STACK/PXE_CMD?
            'IMAGE_TRUST_CMD',
            'PCI_CMD',
            'PARAM_CMD',
            'NEIGHBOUR_CMD',
            'PING_CMD',
            'CONSOLE_CMD',
            'IPSTAT_CMD',
            'PROFSTAT_CMD',
            'NTP_CMD',
            'CERT_CMD'
            ],
        'disable': [
            # 'CRYPTO_80211_WEP',
            # 'CRYPTO_80211_WPA',
            # 'CRYPTO_80211_WPA2',
            # 'IWMGMT_CMD'
            ]},
    'config/console.h': {
        'enable': [
            'CONSOLE_FRAMEBUFFER'
            ],
        'disable': [
            # Disables would go here.
            ]}}

rootdir = '/opt/builds'
logdir = os.path.join(rootdir, 'logs', 'ipxe')
patchdir = os.path.join(rootdir, 'patches', 'ipxe')
buildroot = os.path.join(rootdir, 'ipxe_build')
srcdir = os.path.join(buildroot, 'src')
configdir = os.path.join(rootdir, 'configs', 'ipxe')
destdir = os.path.join(rootdir, 'built', 'ipxe')
repo = git.Repo(buildroot)
master = repo.branches.master
remote = repo.remote('origin')

rom_suffixes = ('rom', 'mrom')
rom_types = ('rtl8139', '8086100e', '80861209', '10500940', '10222000', '10ec8139', '1af41000', '8086100f',
             '808610d3', '15ad07b0')


def doMonkeypatch(fname, changeset):
    enable_re = None
    disable_re = None
    fpath = os.path.join(srcdir, fname)
    if changeset['enable']:
        enable_re = re.compile((r'^\s*(//#define|#undef)\s+'
                                r'(?P<optname>{0})'
                                r'(?P<comment>\s+/\*.*)?\s*$').format('|'.join(changeset['enable'])))
    if changeset['disable']:
        disable_re = re.compile((r'^(#define|//#undef)\s+'
                                 r'(?P<optname>{0})'
                                 r'(?P<comment>\s+/\*.*)?\s*$').format('|'.join(changeset['disable'])))
    with open(fpath, 'r') as fh:
        configstr = fh.read()
    configlines = configstr.splitlines()
    for idx, line in enumerate(configlines[:]):
        if enable_re:
            r = enable_re.search(line)
            if r:
                configlines[idx] = '#define {0}{1}'.format(r.group('optname'), r.group('comment'))
        if disable_re:
            r = disable_re.search(line)
            if r:
                configlines[idx] = '#undef {0}{1}'.format(r.group('optname'), r.group('comment'))
    with open(fpath, 'w') as fh:
        fh.write('\n'.join(configlines))
        fh.write('\n')
    return()

def main():
    # Cleanup the repo.
    repo.head.reset(commit = 'HEAD', working_tree = True)
    if repo.active_branch != master:
        master.checkout()
    repo.head.reset(commit = 'HEAD', working_tree = True)
    repo.git.clean('-xdf')
    try:
        remote.pull()
    except BrokenPipeError:
        pass
    # Patch
    for p in patches:
        with open(os.path.join(patchdir, p), 'rb') as fh:
            patchset = patch_ng.PatchSet(fh)
        patchset.apply(strip = 1, root = buildroot)
    # "Monkeypatch" - sed-like.
    for f, changeset in monkeypatch.items():
        doMonkeypatch(f, changeset)
    # Build. Finally!
    # TODO: ARM support!
    # for d in ('default', 'efi', 'iso', 'legacy', 'roms', 'arm'):
    for d in ('default', 'efi', 'iso', 'legacy', 'roms'):
        dpath = os.path.join(destdir, d)
        os.makedirs(dpath, exist_ok = True)
    os.makedirs(logdir, exist_ok = True)
    os.chdir(srcdir)
    ## Base files
    # TODO: ARM support!
    # TODO: efi-sb support (secureboot)!
    #       http://ipxe.org/appnote/buildtargets
    #       http://ipxe.org/appnote/buildtargets#special_targets
    ## BOOTSTRAP ##
    with open(os.path.join(logdir, 'all_bootstrap.stderr'), 'wb') as stderr, \
            open(os.path.join(logdir, 'all_bootstrap.stdout'), 'wb') as stdout:
        subprocess.run(['make',
                        'all',
                        'EMBED={0}'.format(os.path.join(configdir, 'bootstrap.ipxe'))],
                       stdout = stdout,
                       stderr = stderr)
    for fsrc, fdest in (('undionly.kpxe', 'legacy/bootstrap_{f}'),
                        ('ipxe.iso', 'iso/bootstrap_bios.iso'),
                        ('ipxe.usb', 'iso/bootstrap_pxe_usb.img'),
                        ('ipxe.dsk', 'iso/bootstrap_pxe_floppy.img')):
        srcpath = os.path.join(srcdir, 'bin', fsrc)
        fname = os.path.basename(srcpath)
        destpath = os.path.join(destdir, fdest.format(f = fname))
        shutil.copy2(srcpath, destpath)
    for rom in rom_types:
        for s in rom_suffixes:
            fname = '{0}.{1}'.format(rom, s)
            fpath = os.path.join(srcdir, 'bin', fname)
            if os.path.isfile(fpath):
                shutil.copy2(fpath,
                             os.path.join(destdir, 'roms', 'bootstrap_{0}'.format(fname)))
                # http://ipxe.org/howto/romburning
                # https://libvirt.org/formatdomain.html#elementsNICSROM
                if rom == '1af41000':
                    os.symlink(fpath,
                               os.path.join(destdir, 'roms', 'bootstrap_virtio.rom'))
    ## EMBEDDED ##
    with open(os.path.join(logdir, 'all.stderr'), 'wb') as stderr, \
            open(os.path.join(logdir, 'all.stdout'), 'wb') as stdout:
        subprocess.run(['make',
                        'all',
                        'EMBED={0}'.format(os.path.join(configdir, 'chain-default.ipxe'))],
                        stdout = stdout,
                        stderr = stderr)
    for fsrc, fdest in (('undionly.kpxe', 'legacy/{f}'),
                        ('ipxe.iso', 'iso/bios.iso'),
                        ('ipxe.usb', 'iso/pxe_usb.img'),
                        ('ipxe.dsk', 'iso/pxe_floppy.img')):
        srcpath = os.path.join(srcdir, 'bin', fsrc)
        fname = os.path.basename(srcpath)
        destpath = os.path.join(destdir, fdest.format(f = fname))
        shutil.copy2(srcpath, destpath)
    for rom in rom_types:
        for s in rom_suffixes:
            fname = '{0}.{1}'.format(rom, s)
            fpath = os.path.join(srcdir, 'bin', fname)
            if os.path.isfile(fpath):
                shutil.copy2(fpath,
                             os.path.join(destdir, 'roms', fname))
                if rom == '1af41000':
                    os.symlink(fpath,
                               os.path.join(destdir, 'roms', 'virtio.rom'))
    # DOS/MBR sector/HDD img.
    ## BOOTSTRAP ##
    with open(os.path.join(logdir, 'bootstrap_hdd.stderr'), 'wb') as stderr, \
            open(os.path.join(logdir, 'bootstrap_hdd.stdout'), 'wb') as stdout:
        subprocess.run(['make',
                        'bin/ipxe.hd',
                        'EMBED={0}'.format(os.path.join(configdir, 'bootstrap.ipxe'))],
                       stdout = stdout,
                       stderr = stderr)
    shutil.copy2(os.path.join(srcdir, 'bin/ipxe.hd'),
                 os.path.join(destdir, 'iso', 'bootstrap_legacy_mbr.img'))
    ## EMBEDDED ##
    with open(os.path.join(logdir, 'hdd.stderr'), 'wb') as stderr, \
            open(os.path.join(logdir, 'hdd.stdout'), 'wb') as stdout:
        subprocess.run(['make',
                        'bin/ipxe.hd',
                        'EMBED={0}'.format(os.path.join(configdir, 'chain-default.ipxe'))],
                        stdout = stdout,
                        stderr = stderr)
    shutil.copy2(os.path.join(srcdir, 'bin/ipxe.hd'),
                 os.path.join(destdir, 'iso', 'legacy_mbr.img'))
    # PXE loaders
    ## BOOTSTRAP ##
    with open(os.path.join(logdir, 'bootstrap_loader.stderr'), 'wb') as stderr, \
            open(os.path.join(logdir, 'bootstrap_loader.stdout'), 'wb') as stdout:
        subprocess.run(['make',
                        'bin/ipxe.pxe',
                        'EMBED={0}'.format(os.path.join(configdir, 'bootstrap.ipxe'))],
                       stdout = stdout,
                       stderr = stderr)
    os.rename(os.path.join(srcdir, 'bin/ipxe.pxe'),
              os.path.join(destdir, 'default', 'bootstrap_loader.pxe'))
    ## EMBEDDED ##
    with open(os.path.join(logdir, 'loader.stderr'), 'wb') as stderr, \
            open(os.path.join(logdir, 'loader.stdout'), 'wb') as stdout:
        subprocess.run(['make',
                        'bin/ipxe.pxe',
                        'EMBED={0}'.format(os.path.join(configdir, 'chain-default.ipxe'))],
                       stdout = stdout,
                       stderr = stderr)
    os.rename(os.path.join(srcdir, 'bin/ipxe.pxe'),
              os.path.join(destdir, 'default', 'loader.pxe'))
    # EFI binaries and ISO images
    # These have to be done grouped because the eiso stuff APPARENTLY doesn't parse EMBED or lack thereof correctly?
    ## BOOTSTRAP ##
    ### EFI ###
    with open(os.path.join(logdir, 'bootstrap_efi.stderr'), 'wb') as stderr, \
            open(os.path.join(logdir, 'bootstrap_efi.stdout'), 'wb') as stdout:
        subprocess.run(['make',
                        'bin-i386-efi/ipxe.efi',
                        'bin-x86_64-efi/ipxe.efi',
                        'EMBED={0}'.format(os.path.join(configdir, 'bootstrap.ipxe'))],
                       stdout = stdout,
                       stderr = stderr)
    shutil.copy2(os.path.join(srcdir, 'bin-i386-efi/ipxe.efi'),
                 os.path.join(destdir, 'efi', 'bootstrap_32.efi'))
    shutil.copy2(os.path.join(srcdir, 'bin-x86_64-efi/ipxe.efi'),
                 os.path.join(destdir, 'efi', 'bootstrap_64.efi'))
    ### UEFI ISO ###
    with open(os.path.join(logdir, 'bootstrap_iso.stderr'), 'wb') as stderr, \
            open(os.path.join(logdir, 'bootstrap_iso.stdout'), 'wb') as stdout:
        subprocess.run(['make',
                        'bin/ipxe.liso',
                        'bin/ipxe.eiso',
                        'EMBED={0}'.format(os.path.join(configdir, 'bootstrap.ipxe'))],
                       stdout = stdout,
                       stderr = stderr)
    os.rename(os.path.join(srcdir, 'bin/ipxe.liso'),
              os.path.join(destdir, 'iso', 'bootstrap_legacy.iso'))
    os.rename(os.path.join(srcdir, 'bin/ipxe.eiso'),
              os.path.join(destdir, 'iso', 'bootstrap_uefi.iso'))
    ## EMBEDDED ##
    ### EFI ###
    with open(os.path.join(logdir, 'efi.stderr'), 'wb') as stderr, \
            open(os.path.join(logdir, 'efi.stdout'), 'wb') as stdout:
        subprocess.run(['make',
                        'bin-i386-efi/ipxe.efi',
                        'bin-x86_64-efi/ipxe.efi',
                        'EMBED={0}'.format(os.path.join(configdir, 'chain-default.ipxe'))],
                       stdout = stdout,
                       stderr = stderr)
    shutil.copy2(os.path.join(srcdir, 'bin-i386-efi/ipxe.efi'),
                 os.path.join(destdir, 'efi', '32.efi'))
    shutil.copy2(os.path.join(srcdir, 'bin-x86_64-efi/ipxe.efi'),
                 os.path.join(destdir, 'efi', '64.efi'))
    ### UEFI ISO ###
    with open(os.path.join(logdir, 'iso.stderr'), 'wb') as stderr, \
            open(os.path.join(logdir, 'iso.stdout'), 'wb') as stdout:
        subprocess.run(['make',
                        'bin/ipxe.liso',
                        'bin/ipxe.eiso',
                        'EMBED={0}'.format(os.path.join(configdir, 'chain-default.ipxe'))],
                       stdout = stdout,
                       stderr = stderr)
    os.rename(os.path.join(srcdir, 'bin/ipxe.liso'),
              os.path.join(destdir, 'iso', 'legacy.iso'))
    os.rename(os.path.join(srcdir, 'bin/ipxe.eiso'),
              os.path.join(destdir, 'iso', 'uefi.iso'))
    return()

if __name__ == '__main__':
    main()