#!/usr/bin/env python3
#
# Copyright (C) 2022-2024 VyOS maintainers and contributors
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 or later as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# File: build-vyos-image
# Purpose: builds VyOS images using a fork of Debian's live-build tool

# Import Python's standard library modules
import re
import os
import sys
import copy
import uuid
import glob
import json
import shutil
import argparse
import datetime
import functools
import string

## Check if the script is running wirh root permissions
## Since live-build requires privileged calls such as chroot(),
## there's no real way around it.
if os.getuid() != 0:
    print("E: this script requires root privileges")
    sys.exit(1)

# Import third-party modules
try:
    import tomli
    import jinja2
    import git
    import psutil
except ModuleNotFoundError as e:
    print(f"E: Cannot load required library {e}")
    print("E: Please make sure the following Python3 modules are installed: tomli jinja2 git psutil")
    sys.exit(1)

# Import local defaults
import defaults

## Load the file with default build configuration options
try:
    with open(defaults.DEFAULTS_FILE, 'rb') as f:
        build_defaults = tomli.load(f)
except Exception as e:
    print("E: Failed to open the defaults file {0}: {1}".format(defaults.DEFAULTS_FILE, e))
    sys.exit(1)

# Checkout vyos-1x under build directory
try:
    branch_name = build_defaults['vyos_branch']
    url_vyos_1x = 'https://github.com/vyos/vyos-1x'
    path_vyos_1x = os.path.join(defaults.BUILD_DIR, 'vyos-1x')
    try:
        repo_vyos_1x = git.Repo.clone_from(url_vyos_1x, path_vyos_1x, no_checkout=True)
    except git.GitCommandError:
        if os.path.exists(path_vyos_1x):
            try:
                repo_vyos_1x = git.Repo(path_vyos_1x)
            except git.GitError:
                print(f'E: Corrupted vyos-1x git repo: {path_vyos_1x}; remove')
                sys.exit(1)
        else:
            raise
    # alternatively, pass commit hash or tag as arg:
    repo_vyos_1x.git.checkout(branch_name)
except Exception as e:
    print(f'E: Could not retrieve vyos-1x from branch {branch_name}: {repr(e)}')
    sys.exit(1)

# Add the vyos-1x directory to the Python path so that
# we can import modules from it.
VYOS1X_DIR = os.path.join(os.getcwd(), defaults.BUILD_DIR, 'vyos-1x/python')
if not os.path.exists(VYOS1X_DIR):
    print("E: vyos-1x subdirectory does not exist, did git checkout fail?")
    sys.exit(1)
else:
    sys.path.append(VYOS1X_DIR)

# Import local modules from scripts/image-build
# They rely on modules from vyos-1x
import utils
import raw_image

from utils import cmd

## Check if there are missing build dependencies
deps = {
    'packages': [
       'sudo',
       'make',
       'live-build',
       'pbuilder',
       'devscripts',
       'python3-pystache',
       'python3-git',
       'qemu-utils',
       'gdisk',
       'kpartx',
       'dosfstools'
    ],
   'binaries': []
}

print("I: Checking if packages required for VyOS image build are installed")
try:
    utils.check_system_dependencies(deps)
except OSError as e:
    print(f"E: {e}")
    sys.exit(1)

# argparse converts hyphens to underscores,
# so for lookups in the original options hash we have to convert them back
def field_to_option(s):
    return re.sub(r'_', '-', s)

def get_validator(optdict, name):
    try:
        return optdict[name][1]
    except KeyError:
        return None

def merge_defaults(source, defaults={}, skip_none=False):
    """ Merge a dict with values from a defaults dict.
    Merging logic is as follows:
      Sub-dicts are merged.
      List values are combined.
      Scalar values are set to those from the source dict,
      if they exist there.
    """
    from copy import deepcopy
    tmp = deepcopy(defaults)

    for key, value in source.items():
        if key not in tmp:
            tmp[key] = value
        elif isinstance(source[key], dict):
            tmp[key] = merge_defaults(source[key], tmp[key])
        elif isinstance(source[key], list):
            tmp[key] = source[key] + tmp[key]
        elif not skip_none or source[key] is not None:
            tmp[key] = source[key]

    return tmp

def has_nonempty_key(config, key):
    if key in config:
        if config[key]:
            return True
    return False

def make_toml_path(dir, file_basename):
    return os.path.join(dir, file_basename + ".toml")


if __name__ == "__main__":
    ## Get a list of available build flavors
    flavor_dir_env = os.getenv("VYOS_BUILD_FLAVORS_DIR")
    if flavor_dir_env:
        flavor_dir = flavor_dir_env
    else:
        flavor_dir = defaults.BUILD_FLAVORS_DIR

    print(f"I: using build flavors directory {flavor_dir}")
    build_flavors = [f[0] for f in map(os.path.splitext, os.listdir(flavor_dir)) if (f[1] == ".toml")]

    ## Set up the option parser
    ## XXX: It uses values from the default configuration file for its option defaults,
    ## which is why it's defined after loading the defaults.toml file data.

    # Options dict format:
    # '$option_name_without_leading_dashes': { ('$help_string', $default_value_generator_thunk, $value_checker_thunk) }
    options = {
       'architecture': ('Image target architecture (amd64 or arm64)', lambda x: x in ['amd64', 'arm64', None]),
       'build-by': ('Builder identifier (e.g. jrandomhacker@example.net)', None),
       'debian-mirror': ('Debian repository mirror', None),
       'debian-security-mirror': ('Debian security updates mirror', None),
       'pbuilder-debian-mirror': ('Debian repository mirror for pbuilder env bootstrap', None),
       'vyos-mirror': ('VyOS package mirror', None),
       'build-type': ('Build type, release or development', lambda x: x in ['release', 'development']),
       'version': ('Version number (release builds only)', None),
       'build-comment': ('Optional build comment', None)
    }

    # Create the option parser
    parser = argparse.ArgumentParser()
    for k, v in options.items():
        help_string = v[0]
        parser.add_argument('--' + k, type=str, help=help_string)

    # Debug options
    parser.add_argument('--debug', help='Enable debug output', action='store_true')
    parser.add_argument('--dry-run', help='Check build configuration and exit', action='store_true')

    # Custom APT entry and APT key options can be used multiple times
    parser.add_argument('--custom-apt-entry', help="Custom APT entry", action='append')
    parser.add_argument('--custom-apt-key', help="Custom APT key file", action='append')
    parser.add_argument('--custom-package', help="Custom package to install from repositories", action='append')

    # Options relevant for non-ISO format flavors
    parser.add_argument('--reuse-iso', help='Use an existing ISO file to build additional image formats', type=str, action='store')
    parser.add_argument('--disk-size', help='Disk size for non-ISO image formats', type=int, action='store')

    # Build flavor is a positional argument
    parser.add_argument('build_flavor', help='Build flavor', nargs='?', action='store')

    args = vars(parser.parse_args())

    debug = args["debug"]

    # Validate options
    for k, v in args.items():
        key = field_to_option(k)
        func = get_validator(options, k)
        if func is not None:
            if not func(v):
                print("E: {v} is not a valid value for --{o} option".format(o=key, v=v))
                sys.exit(1)

    if not args["build_flavor"]:
        print("E: Build flavor is not specified!")
        print("E: For example, to build the generic ISO, run {} iso".format(sys.argv[0]))
        print("Available build flavors:\n")
        print("\n".join(build_flavors))
        sys.exit(1)

    ## Try to get correct architecture and build type from build flavor and CLI arguments
    pre_build_config = copy.deepcopy(build_defaults)

    flavor_config = {}
    build_flavor = args["build_flavor"]
    try:
        toml_flavor_file = make_toml_path(flavor_dir, args["build_flavor"])
        with open(toml_flavor_file, 'rb') as f:
            flavor_config = tomli.load(f)
            pre_build_config = merge_defaults(flavor_config, defaults=pre_build_config)
    except FileNotFoundError:
        print(f"E: Flavor '{build_flavor}' does not exist")
        sys.exit(1)
    except tomli.TOMLDecodeError as e:
        print(f"E: Failed to parse TOML file for flavor '{build_flavor}': {e}")
        sys.exit(1)

    ## Combine configs args > flavor > defaults
    pre_build_config = merge_defaults(args, defaults=pre_build_config, skip_none=True)

    # Some fixup for mirror settings.
    # The idea is: if --debian-mirror is specified but --pbuilder-debian-mirror is not,
    # use the --debian-mirror value for both lb and pbuilder bootstrap
    if pre_build_config['debian_mirror'] is None or pre_build_config['debian_security_mirror'] is None:
        print("E: debian_mirror and debian_security_mirror cannot be empty")
        sys.exit(1)

    if pre_build_config['pbuilder_debian_mirror'] is None:
        args['pbuilder_debian_mirror'] = pre_build_config['pbuilder_debian_mirror'] = pre_build_config['debian_mirror']

    # Version can only be set for release builds,
    # for dev builds it hardly makes any sense
    if pre_build_config['build_type'] == 'development':
        if args['version'] is not None:
            print("E: Version can only be set for release builds")
            print("Use --build-type=release option if you want to set version number")
            sys.exit(1)

    # Validate characters in version name
    if 'version' in args and args['version'] != None:
        allowed = string.ascii_letters + string.digits + '.' + '-' + '+'
        if not set(args['version']) <= set(allowed):
            print(f'Version contained illegal character(s), allowed: {allowed}')
            sys.exit(1)

    ## Inject some useful hardcoded options
    args['build_dir'] = defaults.BUILD_DIR
    args['pbuilder_config'] = os.path.join(defaults.BUILD_DIR, defaults.PBUILDER_CONFIG)

    ## Add hardcoded defaults
    build_config = merge_defaults(defaults.HARDCODED_BUILD, defaults={})

    ## Combine the arguments with non-configurable defaults
    build_config = merge_defaults(build_defaults, defaults=build_config)

    ## Load correct mix-ins
    with open(make_toml_path(defaults.BUILD_TYPES_DIR, pre_build_config["build_type"]), 'rb') as f:
        build_type_config = tomli.load(f)
        build_config = merge_defaults(build_type_config, defaults=build_config)

    with open(make_toml_path(defaults.BUILD_ARCHES_DIR, pre_build_config["architecture"]), 'rb') as f:
        build_arch_config = tomli.load(f)
        build_config = merge_defaults(build_arch_config, defaults=build_config)

    ## Override with flavor and then CLI arguments
    build_config = merge_defaults(flavor_config, defaults=build_config)
    build_config = merge_defaults(args, defaults=build_config, skip_none=True)

    ## Rename and merge some fields for simplicity
    ## E.g. --custom-packages is for the user, but internally
    ## it's added to the same package list as everything else
    if has_nonempty_key(build_config, "custom_package"):
        build_config["packages"] += build_config["custom_package"]
        del build_config["custom_package"]

    ## Add architecture-dependent packages from the flavor
    if has_nonempty_key(build_config, "architectures"):
        arch = build_config["architecture"]
        if arch in build_config["architectures"]:
            if has_nonempty_key(build_config["architectures"][arch], "packages"):
                build_config["packages"] += build_config["architectures"][arch]["packages"]

    ## Check if image format is specified,
    ## else we have no idea what we are actually supposed to build.
    if not has_nonempty_key(build_config, "image_format"):
        print("E: image format is not specified in the build flavor file")
        sys.exit(1)

    ## Add default boot settings if needed
    if "boot_settings" not in build_config:
        build_config["boot_settings"] = defaults.boot_settings
    else:
        build_config["boot_settings"] = merge_defaults(build_config["boot_settings"], defaults=defaults.boot_settings)

    ## Convert the image_format field to a single-item list if it's a scalar
    ## (like `image_format = "iso"`)
    if type(build_config["image_format"]) != list:
        build_config["image_format"] = [ build_config["image_format"] ]

    ## Dump the complete config if the user enabled debug mode
    if debug:
        import json
        print("D: Effective build config:\n")
        print(json.dumps(build_config, indent=4))

    ## Clean up the old build config and set up a fresh copy
    lb_config_dir = os.path.join(defaults.BUILD_DIR, defaults.LB_CONFIG_DIR)
    shutil.rmtree(lb_config_dir, ignore_errors=True)
    shutil.copytree("data/live-build-config/", lb_config_dir)
    os.makedirs(lb_config_dir, exist_ok=True)

    # Switch to the build directory, this is crucial for the live-build work
    # because the efective build config files etc. are there.
    #
    # All directory paths from this point must be relative to BUILD_DIR,
    # not to the vyos-build repository root.
    os.chdir(defaults.BUILD_DIR)

    iso_file = None

    if build_config["reuse_iso"]:
        iso_file = build_config["reuse_iso"]
    else:
        ## Create the version file

        # Create a build timestamp
        now = datetime.datetime.today()
        build_timestamp = now.strftime("%Y%m%d%H%M")

        # FIXME: use aware rather than naive object
        build_date = now.strftime("%a %d %b %Y %H:%M UTC")

        # Assign a (hopefully) unique identifier to the build (UUID)
        build_uuid = str(uuid.uuid4())

        # Initialize Git object from our repository
        try:
            repo = git.Repo('.', search_parent_directories=True)
            # Retrieve the Git commit ID of the repository, 14 charaters will be sufficient
            build_git = repo.head.object.hexsha[:14]
            # If somone played around with the source tree and the build is "dirty", mark it
            if repo.is_dirty():
                build_git += "-dirty"

            # Retrieve git branch name or current tag
            # Building a tagged release might leave us checking out a git tag that is not the tip of a named branch (detached HEAD)
            # Check if the current HEAD is associated with a tag and use its name instead of an unavailable branch name.
            git_branch = next((tag.name for tag in repo.tags if tag.commit == repo.head.commit), None)
            if git_branch is None:
                git_branch = repo.active_branch.name
        except Exception as e:
            print(f'W: Could not retrieve information from git: {repr(e)}')
            build_git = ""
            git_branch = ""

        # Create the build version string
        if build_config['build_type'] == 'development':
            try:
                if not git_branch:
                    raise ValueError("git branch could not be determined")

                # Load the branch to version mapping file
                with open('../data/versions') as f:
                    version_mapping = json.load(f)

                branch_version = version_mapping[git_branch]

                version = "{0}-rolling-{1}".format(branch_version, build_timestamp)
            except Exception as e:
                print("W: Could not build a version string specific to git branch, falling back to default: {0}".format(str(e)))
                version = "999.{0}".format(build_timestamp)
        else:
            # Release build, use the version from ./configure arguments
            version = build_config['version']

        if build_config['build_type'] == 'development':
            lts_build = False
        else:
            lts_build = True

        version_data = {
            'version': version,
            'flavor': build_config["build_flavor"],
            'built_by': build_config['build_by'],
            'built_on': build_date,
            'build_uuid': build_uuid,
            'build_git': build_git,
            'build_branch': git_branch,
            'release_train': build_config['release_train'],
            'architecture': build_config['architecture'],
            'lts_build': lts_build,
            'build_comment': build_config['build_comment'],
            'bugtracker_url': build_config['bugtracker_url'],
            'documentation_url': build_config['documentation_url'],
            'project_news_url': build_config['project_news_url'],
            'support_url': build_config['support_url']
        }

        # Multi line strings needs to be un-indented to not have leading
        # whitespaces in the resulting file
        os_release = f"""
    PRETTY_NAME="VyOS {version} ({build_config['release_train']})"
    NAME="VyOS"
    VERSION_ID="{version}"
    VERSION="{version} ({build_config['release_train']})"
    VERSION_CODENAME={build_defaults['debian_distribution']}
    ID=vyos
    BUILD_ID="{build_git}"
    HOME_URL="{build_defaults['website_url']}"
    SUPPORT_URL="{build_defaults['support_url']}"
    BUG_REPORT_URL="{build_defaults['bugtracker_url']}"
    DOCUMENTATION_URL="{build_config['documentation_url']}"
        """

        # Reminder: all paths relative to the build dir, not to the repository root
        chroot_includes_dir = defaults.CHROOT_INCLUDES_DIR
        binary_includes_dir = defaults.BINARY_INCLUDES_DIR
        vyos_data_dir = os.path.join(chroot_includes_dir, "usr/share/vyos")
        os.makedirs(vyos_data_dir, exist_ok=True)
        with open(os.path.join(vyos_data_dir, 'version.json'), 'w') as f:
            json.dump(version_data, f)
        with open(os.path.join(binary_includes_dir, 'version.json'), 'w') as f:
            json.dump(version_data, f)

        # For backwards compatibility with 'add system image' script from older versions
        # we need a file in the old format so that script can find out the version of the image
        # for upgrade
        os.makedirs(os.path.join(chroot_includes_dir, 'opt/vyatta/etc/'), exist_ok=True)
        with open(os.path.join(chroot_includes_dir, 'opt/vyatta/etc/version'), 'w') as f:
            print("Version: {0}".format(version), file=f)

        # Define variables that influence to welcome message on boot
        os.makedirs(os.path.join(chroot_includes_dir, 'usr/lib/'), exist_ok=True)
        with open(os.path.join(chroot_includes_dir, 'usr/lib/os-release'), 'w') as f:
            print(os_release, file=f)

        ## Clean up earlier build state and artifacts
        print("I: Cleaning the build workspace")
        cmd("lb clean")
        #iter(lambda p: shutil.rmtree(p, ignore_errors=True),
        #  ['config/binary', 'config/bootstrap', 'config/chroot', 'config/common', 'config/source'])
        artifacts = functools.reduce(
          lambda x, y: x + y,
          map(glob.glob, ['*.iso', '*.raw', '*.img', '*.xz', '*.ova', '*.ovf']))
        iter(os.remove, artifacts)

        ## Create the target ISO file path
        iso_file = f"vyos-{version_data['version']}-{build_config['build_flavor']}-{build_config['architecture']}.iso"

        ## Create live-build configuration files

        # Add the additional repositories to package lists
        print("I: Setting up additional APT entries")
        vyos_repo_entry = "deb {vyos_mirror} {vyos_branch} main\n".format(**build_config)

        apt_file = defaults.VYOS_REPO_FILE

        if debug:
            print(f"D: Adding these entries to {apt_file}:")
            print("\t", vyos_repo_entry)

        with open(apt_file, 'w') as f:
            f.write(vyos_repo_entry)

        # Add custom APT entries
        if build_config.get('additional_repositories', False):
            build_config['custom_apt_entry'] += build_config['additional_repositories']

        if build_config.get('custom_apt_entry', False):
            custom_apt_file = defaults.CUSTOM_REPO_FILE
            entries = "\n".join(build_config['custom_apt_entry'])
            if debug:
                print("D: Adding custom APT entries:")
                print(entries)
            with open(custom_apt_file, 'w') as f:
                f.write(entries)
                f.write("\n")

        # Add custom APT keys
        if has_nonempty_key(build_config, 'custom_apt_key'):
            key_dir = defaults.ARCHIVES_DIR
            for k in build_config['custom_apt_key']:
                dst_name = '{0}.key.chroot'.format(os.path.basename(k))
                shutil.copy(k, os.path.join(key_dir, dst_name))

        # Add custom packages
        if has_nonempty_key(build_config, 'packages'):
            package_list_file = defaults.PACKAGE_LIST_FILE
            packages = "\n".join(build_config['packages'])
            with open (package_list_file, 'w') as f:
                f.write(packages)

        ## Create includes
        if has_nonempty_key(build_config, "includes_chroot"):
            for i in build_config["includes_chroot"]:
                file_path = os.path.join(chroot_includes_dir, i["path"])
                if debug:
                    print(f"D: Creating chroot include file: {file_path}")
                os.makedirs(os.path.dirname(file_path), exist_ok=True)
                with open(file_path, 'w') as f:
                    f.write(i["data"])

        ## Create the default config
        ## Technically it's just another includes.chroot entry,
        ## but it's special enough to warrant making it easier for flavor writers
        if has_nonempty_key(build_config, "default_config"):
            file_path = os.path.join(chroot_includes_dir, "opt/vyatta/etc/config.boot.default")
            os.makedirs(os.path.dirname(file_path), exist_ok=True)
            with open(file_path, 'w') as f:
                f.write(build_config["default_config"])

        ## Initialize build manifest
        manifest = {
            'build_config' : build_config,
            'artifacts' : [iso_file],
            'pre_build_config' : pre_build_config
        }

        ## Configure live-build
        lb_config_tmpl = jinja2.Template("""
        lb config noauto \
                --apt-indices false \
                --apt-options "--yes -oAPT::Get::allow-downgrades=true" \
                --apt-recommends false \
                --architecture {{architecture}} \
                --archive-areas {{debian_archive_areas}} \
                --backports true \
                --binary-image iso-hybrid \
                --bootappend-live "boot=live components hostname=vyos username=live nopersistence noautologin nonetworking union=overlay console=ttyS0,115200 console=tty0 net.ifnames=0 biosdevname=0" \
                --bootappend-live-failsafe "live components memtest noapic noapm nodma nomce nolapic nomodeset nosmp nosplash vga=normal console=ttyS0,115200 console=tty0 net.ifnames=0 biosdevname=0" \
                --bootloaders {{bootloaders}} \
                --checksums 'sha256 md5' \
                --chroot-squashfs-compression-type "{{squashfs_compression_type}}" \
                --debian-installer none \
                --debootstrap-options "--variant=minbase --exclude=isc-dhcp-client,isc-dhcp-common,ifupdown --include=apt-utils,ca-certificates,gnupg2" \
                --distribution {{debian_distribution}} \
                --firmware-binary false \
                --firmware-chroot false \
                --iso-application "VyOS" \
                --iso-publisher "{{build_by}}" \
                --iso-volume "VyOS" \
                --linux-flavours {{kernel_flavor}} \
                --linux-packages linux-image-{{kernel_version}} \
                --mirror-binary {{debian_mirror}} \
                --mirror-binary-security {{debian_security_mirror}} \
                --mirror-bootstrap {{debian_mirror}} \
                --mirror-chroot {{debian_mirror}} \
                --mirror-chroot-security {{debian_security_mirror}} \
                --security true \
                --updates true
                "${@}"
        """)

        lb_config_command = lb_config_tmpl.render(build_config)

        ## Pin release for VyOS packages
        apt_pin = f"""Package: *
Pin: release n={build_config['release_train']}
Pin-Priority: 600
"""

        with open(defaults.VYOS_PIN_FILE, 'w') as f:
            f.write(apt_pin)

        print("I: Configuring live-build")

        if debug:
            print("D: live-build configuration command")
            print(lb_config_command)

        cmd(lb_config_command)

        ## In dry-run mode, exit at this point
        if build_config["dry_run"]:
            print("I: dry-run, not starting image build")
            sys.exit(0)

        ## Add local packages
        local_packages = glob.glob('../packages/*.deb')
        if local_packages:
            for f in local_packages:
                shutil.copy(f, os.path.join(defaults.LOCAL_PACKAGES_PATH, os.path.basename(f)))

        ## Build the image
        print("I: Starting image build")
        if debug:
            print("D: It's not like I'm building this specially for you or anything!")
        cmd("lb build 2>&1")

        # Copy the image
        shutil.copy("live-image-{0}.hybrid.iso".format(build_config["architecture"]), iso_file)

    # Build additional flavors from the ISO,
    # if the flavor calls for them
    if build_config["image_format"] != ["iso"]:
        raw_image = raw_image.create_raw_image(build_config, iso_file, "tmp/")
        manifest['artifacts'].append(raw_image)

        if has_nonempty_key(build_config, "post_build_hook"):
            # Some flavors require special procedures that aren't covered by qemu-img
            # (most notably, the VMware OVA that requires a custom tool to make and sign the image).
            # For those cases, we support running a post-build hook on the raw image.
            # The image_format field should be 'raw' if a post-build hook is used.
            hook_path = build_config["post_build_hook"]
            cmd(f"{hook_path} {raw_image}")
        else:
            # Most other formats, thankfully, can be produced with just `qemu-img convert`
            other_formats = filter(lambda x: x not in ["iso", "raw"], build_config["image_format"])
            for f in other_formats:
                image_ext = build_config.get("image_ext", f)
                image_opts = build_config.get("image_opts", "")
                target = f"{os.path.splitext(raw_image)[0]}.{image_ext}"
                print(f"I: Building {f} file {target}")
                cmd(f"qemu-img convert -f raw -O {f} {image_opts} {raw_image} {target}")
                manifest['artifacts'].append(target)

    with open('manifest.json', 'w') as f:
        f.write(json.dumps(manifest))