#!/usr/bin/env python3
import argparse
import http.client
import json
import os
import socket
import subprocess
import sys

try:
    import lxc
except ImportError:
    print("You must have python3-lxc installed for this script to work.")
    sys.exit(1)


# Whitelist of keys we either need to check or allow setting in LXD. The latter
# is strictly only true for 'lxc.aa_profile'.
keys_to_check = [
    'lxc.pts',
    # 'lxc.tty',
    # 'lxc.devttydir',
    # 'lxc.kmsg',
    'lxc.aa_profile',
    # 'lxc.cgroup.',
    'lxc.loglevel',
    # 'lxc.logfile',
    'lxc.mount.auto',
    'lxc.mount',
    # 'lxc.rootfs.mount',
    # 'lxc.rootfs.options',
    # 'lxc.pivotdir',
    # 'lxc.hook.pre-start',
    # 'lxc.hook.pre-mount',
    # 'lxc.hook.mount',
    # 'lxc.hook.autodev',
    # 'lxc.hook.start',
    # 'lxc.hook.stop',
    # 'lxc.hook.post-stop',
    # 'lxc.hook.clone',
    # 'lxc.hook.destroy',
    # 'lxc.hook',
    'lxc.network.type',
    'lxc.network.flags',
    'lxc.network.link',
    'lxc.network.name',
    'lxc.network.macvlan.mode',
    'lxc.network.veth.pair',
    # 'lxc.network.script.up',
    # 'lxc.network.script.down',
    'lxc.network.hwaddr',
    'lxc.network.mtu',
    # 'lxc.network.vlan.id',
    # 'lxc.network.ipv4.gateway',
    # 'lxc.network.ipv4',
    # 'lxc.network.ipv6.gateway',
    # 'lxc.network.ipv6',
    # 'lxc.network.',
    # 'lxc.network',
    # 'lxc.console.logfile',
    # 'lxc.console',
    'lxc.include',
    'lxc.start.auto',
    'lxc.start.delay',
    'lxc.start.order',
    # 'lxc.monitor.unshare',
    # 'lxc.group',
    'lxc.environment',
    # 'lxc.init_cmd',
    # 'lxc.init_uid',
    # 'lxc.init_gid',
    # 'lxc.ephemeral',
    # 'lxc.syslog',
    # 'lxc.no_new_privs',

    # Additional keys that are either set by this script or are used to report
    # helpful errors to users.
    'lxc.arch',
    'lxc.id_map',
    'lxd.migrated',
    'lxc.rootfs.backend',
    'lxc.rootfs',
    'lxc.utsname',
    'lxc.aa_allow_incomplete',
    'lxc.autodev',
    'lxc.haltsignal',
    'lxc.rebootsignal',
    'lxc.stopsignal',
    'lxc.mount.entry',
    'lxc.cap.drop',
    # 'lxc.cap.keep',
    'lxc.seccomp',
    # 'lxc.se_context',
    ]


# Unix connection to LXD
class UnixHTTPConnection(http.client.HTTPConnection):
    def __init__(self, path):
        http.client.HTTPConnection.__init__(self, 'localhost')
        self.path = path

    def connect(self):
        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        sock.connect(self.path)
        self.sock = sock


# Fetch a config key as a list
def config_get(config, key, default=None):
    result = []
    for line in config:
        fields = line.split("=", 1)
        if fields[0].strip() == key:
            result.append(fields[-1].strip())

    if len(result) == 0:
        return default
    else:
        return result


def config_keys(config):
    keys = []
    for line in config:
        fields = line.split("=", 1)
        cur = fields[0].strip()
        if cur and not cur.startswith("#") and cur.startswith("lxc."):
            keys.append(cur)

    return keys


# Parse a LXC configuration file, called recursively for includes
def config_parse(path):
    config = []
    with open(path, "r") as fd:
        for line in fd:
            line = line.strip()
            key = line.split("=", 1)[0].strip()
            value = line.split("=", 1)[-1].strip()

            # Parse user-added includes
            if key == "lxc.include":
                # Ignore our own default configs
                if value.startswith("/usr/share/lxc/config/"):
                    continue

                if os.path.isfile(value):
                    config += config_parse(value)
                    continue
                elif os.path.isdir(value):
                    for entry in os.listdir(value):
                        if not entry.endswith(".conf"):
                            continue

                        config += config_parse(os.path.join(value, entry))
                    continue
                else:
                    print("Invalid include: %s", line)

            # Expand any fstab
            if key == "lxc.mount":
                if not os.path.exists(value):
                    print("Container fstab file doesn't exist, skipping...")
                    continue

                with open(value, "r") as fd:
                    for line in fd:
                        line = line.strip()
                        if (line and not line.startswith("#") and
                                line.startswith("lxc.")):
                            config.append("lxc.mount.entry = %s" % line)
                continue

            # Proces normal configuration keys
            if line and not line.strip().startswith("#"):
                config.append(line)

    return config


def container_exists(lxd_socket, container_name):
    lxd = UnixHTTPConnection(lxd_socket)
    lxd.request("GET", "/1.0/containers/%s" % container_name)
    if lxd.getresponse().status == 404:
        return False

    return True


def container_create(lxd_socket, args):
    # Define the container
    lxd = UnixHTTPConnection(lxd_socket)
    lxd.request("POST", "/1.0/containers", json.dumps(args))
    r = lxd.getresponse()

    # Decode the response
    resp = json.loads(r.read().decode())
    if resp["type"] == "error":
        raise Exception("Failed to define container: %s" % resp["error"])

    # Wait for result
    lxd = UnixHTTPConnection(lxd_socket)
    lxd.request("GET", "%s/wait" % resp["operation"])
    r = lxd.getresponse()

    # Decode the response
    resp = json.loads(r.read().decode())
    if resp["type"] == "error":
        raise Exception("Failed to define container: %s" % resp["error"])


# Convert a LXC container to a LXD one
def convert_container(lxd_socket, container_name, args):
    print("==> Processing container: %s" % container_name)

    # Load the container
    try:
        container = lxc.Container(container_name, args.lxcpath)
    except Exception:
        print("Invalid container configuration, skipping...")
        return False

    # As some keys can't be queried over the API, parse the config ourselves
    print("Parsing LXC configuration")
    lxc_config = config_parse(container.config_file_name)
    found_keys = config_keys(lxc_config)

    # Generic check for any invalid LXC configuration keys.
    print("Checking for unsupported LXC configuration keys")
    diff = list(set(found_keys) - set(keys_to_check))
    for d in diff:
        if (not d.startswith('lxc.network.') and not
                d.startswith('lxc.cgroup.')):
            print("Found at least one unsupported config key %s: " % d)
            print("Not importing this container, skipping...")
            return False

    if args.debug:
        print("Container configuration:")
        print(" ", end="")
        print("\n ".join(lxc_config))
        print("")

    # Check for keys that have values differing from the LXD defaults.
    print("Checking whether container has already been migrated")
    if config_get(lxc_config, "lxd.migrated"):
        print("Container has already been migrated, skipping...")
        return False

    # Make sure we don't have a conflict
    print("Checking for existing containers")
    if container_exists(lxd_socket, container_name):
        print("Container already exists, skipping...")
        return False

    # Validating lxc.id_map: must be unset.
    print("Validating container mode")
    if config_get(lxc_config, "lxc.id_map"):
        print("Unprivileged containers aren't supported, skipping...")
        return False

    # Validate lxc.utsname
    print("Validating container name")
    value = config_get(lxc_config, "lxc.utsname")
    if value and value[0] != container_name:
        print("Container name doesn't match lxc.utsname, skipping...")
        return False

    # Validate lxc.aa_allow_incomplete: must be set to 0 or unset.
    print("Validating whether incomplete AppArmor support is enabled")
    value = config_get(lxc_config, "lxc.aa_allow_incomplete")
    if value and int(value[0]) != 0:
        print("Container allows incomplete AppArmor support, skipping...")
        return False

    # Validate lxc.autodev: must be set to 1 or unset.
    print("Validating whether mounting a minimal /dev is enabled")
    value = config_get(lxc_config, "lxc.autodev")
    if value and int(value[0]) != 1:
        print("Container doesn't mount a minimal /dev filesystem, skipping...")
        return False

    # Validate lxc.haltsignal: must be unset.
    print("Validating that no custom haltsignal is set")
    value = config_get(lxc_config, "lxc.haltsignal")
    if value:
        print("Container sets custom halt signal, skipping...")
        return False

    # Validate lxc.rebootsignal: must be unset.
    print("Validating that no custom rebootsignal is set")
    value = config_get(lxc_config, "lxc.rebootsignal")
    if value:
        print("Container sets custom reboot signal, skipping...")
        return False

    # Validate lxc.stopsignal: must be unset.
    print("Validating that no custom stopsignal is set")
    value = config_get(lxc_config, "lxc.stopsignal")
    if value:
        print("Container sets custom stop signal, skipping...")
        return False

    # Extract and valid rootfs key
    print("Validating container rootfs")
    value = config_get(lxc_config, "lxc.rootfs")
    if not value:
        print("Invalid container, missing lxc.rootfs key, skipping...")
        return False

    rootfs = value[0]

    if not os.path.exists(rootfs):
        print("Couldn't find the container rootfs '%s', skipping..." % rootfs)
        return False

    # Base config
    config = {}
    config['security.privileged'] = "true"
    devices = {}
    devices['eth0'] = {'type': "none"}

    # Convert network configuration
    print("Processing network configuration")
    try:
        count = len(container.get_config_item("lxc.network"))
    except Exception:
        count = 0

    for i in range(count):
        device = {"type": "nic"}

        # Get the device type
        device["nictype"] = container.get_config_item("lxc.network")[i]

        # Get everything else
        dev = container.network[i]

        # Validate configuration
        if dev.ipv4 or dev.ipv4_gateway:
            print("IPv4 network configuration isn't supported, skipping...")
            return False

        if dev.ipv6 or dev.ipv6_gateway:
            print("IPv6 network configuration isn't supported, skipping...")
            return False

        if dev.script_up or dev.script_down:
            print("Network config scripts aren't supported, skipping...")
            return False

        if device["nictype"] == "none":
            print("\"none\" network mode isn't supported, skipping...")
            return False

        if device["nictype"] == "vlan":
            print("\"vlan\" network mode isn't supported, skipping...")
            return False

        # Convert the configuration
        if dev.hwaddr:
            device['hwaddr'] = dev.hwaddr

        if dev.link:
            device['parent'] = dev.link

        if dev.mtu:
            device['mtu'] = dev.mtu

        if dev.name:
            device['name'] = dev.name

        if dev.veth_pair:
            device['host_name'] = dev.veth_pair

        if device["nictype"] == "veth":
            if "parent" in device:
                device["nictype"] = "bridged"
            else:
                device["nictype"] = "p2p"

        if device["nictype"] == "phys":
            device["nictype"] = "physical"

        if device["nictype"] == "empty":
            continue

        devices['convert_net%d' % i] = device
        count += 1

    # Convert storage configuration
    value = config_get(lxc_config, "lxc.mount.entry", [])
    i = 0
    for entry in value:
        mount = entry.split(" ")
        if len(mount) < 4:
            print("Invalid mount configuration, skipping...")
            return False

        # Ignore mounts that are present in LXD containers by default.
        if mount[0] in ("proc", "sysfs"):
            continue

        device = {'type': "disk"}

        # Deal with read-only mounts
        if "ro" in mount[3].split(","):
            device['readonly'] = "true"

        # Deal with optional mounts
        if "optional" in mount[3].split(","):
            device['optional'] = "true"
        elif not os.path.exists(mount[0]):
            print("Invalid mount configuration, source path doesn't exist.")
            return False

        # Set the source
        device['source'] = mount[0]

        # Figure out the target
        if mount[1][0] != "/":
            device['path'] = "/%s" % mount[1]
        else:
            device['path'] = mount[1].split(rootfs, 1)[-1]

        devices['convert_mount%d' % i] = device
        i += 1

    # Convert environment
    print("Processing environment configuration")
    value = config_get(lxc_config, "lxc.environment", [])
    for env in value:
        entry = env.split("=", 1)
        config['environment.%s' % entry[0].strip()] = entry[-1].strip()

    # Convert auto-start
    print("Processing container boot configuration")
    value = config_get(lxc_config, "lxc.start.auto")
    if value and int(value[0]) > 0:
        config['boot.autostart'] = "true"

    value = config_get(lxc_config, "lxc.start.delay")
    if value and int(value[0]) > 0:
        config['boot.autostart.delay'] = value[0]

    value = config_get(lxc_config, "lxc.start.order")
    if value and int(value[0]) > 0:
        config['boot.autostart.priority'] = value[0]

    # Convert apparmor
    print("Processing container apparmor configuration")
    value = config_get(lxc_config, "lxc.aa_profile")
    if value:
        if value[0] == "lxc-container-default-with-nesting":
            config['security.nesting'] = "true"
        elif value[0] != "lxc-container-default":
            config["raw.lxc"] = "lxc.aa_profile=%s" % value[0]

    # Convert seccomp
    print("Processing container seccomp configuration")
    value = config_get(lxc_config, "lxc.seccomp")
    if value and value[0] != "/usr/share/lxc/config/common.seccomp":
        print("Custom seccomp profiles aren't supported, skipping...")
        return False

    # Convert SELinux
    print("Processing container SELinux configuration")
    value = config_get(lxc_config, "lxc.se_context")
    if value:
        print("Custom SELinux policies aren't supported, skipping...")
        return False

    # Convert capabilities
    print("Processing container capabilities configuration")
    value = config_get(lxc_config, "lxc.cap.drop")
    if value:
        for cap in value:
            # Ignore capabilities that are dropped in LXD containers by default.
            if cap in ("mac_admin", "mac_override", "sys_module", "sys_time"):
                continue
            print("Custom capabilities aren't supported, skipping...")
            return False

    value = config_get(lxc_config, "lxc.cap.keep")
    if value:
        print("Custom capabilities aren't supported, skipping...")
        return False

    # Setup the container creation request
    new = {'name': container_name,
           'source': {'type': 'none'},
           'config': config,
           'devices': devices,
           'profiles': ["default"]}

    # Set the container architecture if set in LXC
    print("Processing container architecture configuration")
    arches = {'i686': "i686",
              'x86_64': "x86_64",
              'armhf': "armv7l",
              'arm64': "aarch64",
              'powerpc': "ppc",
              'powerpc64': "ppc64",
              'ppc64el': "ppc64le",
              's390x': "s390x"}

    arch = None
    try:
        arch = config_get(lxc_config, "lxc.arch", None)

        if arch and arch[0] in arches:
            new['architecture'] = arches[arch[0]]
        else:
            print("Unknown architecture, assuming native.")
    except Exception:
        print("Couldn't find container architecture, assuming native.")

    # Define the container in LXD
    if args.debug:
        print("LXD container config:")
        print(json.dumps(new, indent=True, sort_keys=True))

    if args.dry_run:
        return True

    if container.running:
        print("Only stopped containers can be migrated, skipping...")
        return False

    try:
        print("Creating the container")
        container_create(lxd_socket, new)
    except Exception as e:
        raise
        print("Failed to create the container: %s" % e)
        return False

    # Transfer the filesystem
    lxd_rootfs = os.path.join(args.lxdpath, "containers",
                              container_name, "rootfs")

    if args.move_rootfs:
        if os.path.exists(lxd_rootfs):
            os.rmdir(lxd_rootfs)

        if subprocess.call(["mv", rootfs, lxd_rootfs]) != 0:
            print("Failed to move the container rootfs, skipping...")
            return False

        os.mkdir(rootfs)
    else:
        print("Copying container rootfs")
        if not os.path.exists(lxd_rootfs):
            os.mkdir(lxd_rootfs)

        if subprocess.call(["rsync", "-Aa", "--sparse",
                            "--acls", "--numeric-ids", "--hard-links",
                            "%s/" % rootfs, "%s/" % lxd_rootfs]) != 0:
            print("Failed to transfer the container rootfs, skipping...")
            return False

    # Delete the source
    if args.delete:
        print("Deleting source container")
        container.delete()

    # Mark the container as migrated
    with open(container.config_file_name, "a") as fd:
        fd.write("lxd.migrated=true\n")
    print("Container is ready to use")
    return True


# Argument parsing
parser = argparse.ArgumentParser()
parser.add_argument("--dry-run", action="store_true", default=False,
                    help="Dry run mode")
parser.add_argument("--debug", action="store_true", default=False,
                    help="Print debugging output")
parser.add_argument("--all", action="store_true", default=False,
                    help="Import all containers")
parser.add_argument("--delete", action="store_true", default=False,
                    help="Delete the source container")
parser.add_argument("--move-rootfs", action="store_true", default=False,
                    help="Move the container rootfs rather than copying it")
parser.add_argument("--lxcpath", type=str, default=False,
                    help="Alternate LXC path")
parser.add_argument("--lxdpath", type=str, default="/var/lib/lxd",
                    help="Alternate LXD path")
parser.add_argument(dest='containers', metavar="CONTAINER", type=str,
                    help="Container to import", nargs="*")
args = parser.parse_args()

# Sanity checks
if not os.geteuid() == 0:
    parser.error("You must be root to run this tool")

if (not args.containers and not args.all) or (args.containers and args.all):
    parser.error("You must either pass container names or --all")

# Connect to LXD
lxd_socket = os.path.join(args.lxdpath, "unix.socket")

if not os.path.exists(lxd_socket):
    print("LXD isn't running.")
    sys.exit(1)

# Run migration
results = {}
count = 0
for container_name in lxc.list_containers(config_path=args.lxcpath):
    if args.containers and container_name not in args.containers:
        continue

    if count > 0:
        print("")

    results[container_name] = convert_container(lxd_socket,
                                                container_name, args)
    count += 1

# Print summary
if not results:
    print("No container to migrate")
    sys.exit(0)

print("")
print("==> Migration summary")
for name, result in results.items():
    if result:
        print("%s: SUCCESS" % name)
    else:
        print("%s: FAILURE" % name)

if False in results.values():
    sys.exit(1)
