#!/usr/bin/python3

import argparse
import libvirt
import os
import shlex
import socket
import subprocess
import time
import yaml
from xml.etree import ElementTree


def domain_to_addr(name, uri='qemu:///system'):
    v = libvirt.open('qemu:///system')
    domain = v.lookupByName(name)
    if domain.isActive() != 1:
        raise Exception("VM is not running")

    ifaces = domain.interfaceAddresses(0)
    for iface in ifaces.values():
        return iface['addrs'][0]['addr']

    ifaces = domain.interfaceAddresses(1)
    for iface, data in ifaces.items():
        if iface != 'lo':
            for addr in data['addrs']:
                if addr['type'] == 1 and addr['addr'][0:6] == 'fe80::':
                    continue
                return addr['addr']

    raise Exception("No interfaces found on VM")


def domain_to_cid(name, uri='qemu:///system'):
    v = libvirt.openReadOnly('qemu:///system')
    domain = v.lookupByName(name)
    if domain.isActive() != 1:
        raise Exception("VM is not running")

    et = ElementTree.fromstring(domain.XMLDesc())
    return et.find('devices/vsock/cid').attrib['address']


def scp_if_exists(path, host, user=None, scp_args=['scp']):
    local_path = os.path.expanduser(os.path.join('~', path))
    if os.path.exists(path):
        if user is not None:
            dest = '{0}@{1}:{2}'.format(user, host, path)
        else:
            dest = '{0}:{1}'.format(host, path)

        subprocess.call(scp_args + [local_path, dest])


def xpra_attach(cid, xpra_title):
    return subprocess.call([
        'xpra', 'attach',
        'vsock://{0}:14500'.format(cid),
        '--encoding=rgb',
        '--title={0}'.format(xpra_title),
        '--tray=no',
        '--notifications=no',
        '--username=user',
        '--desktop-scaling=no',
    ])


if __name__ == '__main__':
    default_user = "user"
    default_identity = None

    try:
        import xdg.BaseDirectory
        configpath = xdg.BaseDirectory.load_first_config(
            'appvm/config.yml')
    except ModuleNotFoundError:
        configpath = os.path.expanduser('~/.config/appvm/config.yml')

    if os.path.exists(configpath):
        config = yaml.safe_load(open(configpath))
        default_user = config.get('user', default_user)
        default_identity = config.get('identity', default_identity)

    parser = argparse.ArgumentParser(
        description="Connect to AppVMs using Xpra over SSH")
    parser.add_argument('--cid', required=False,
                        help="vsock CID (required if not using libvirt)")
    parser.add_argument('--user', '-u', '-l', default=default_user,
                        help="SSH username")
    parser.add_argument('--identity', '-i', default=default_identity,
                        help="SSH identity file")
    parser.add_argument('--shell', action='store_true',
                        help="Connect using SSH instead of Xpra")
    parser.add_argument('--skip-libvirt', action='store_true',
                        help="Skip Libvirt (use domain name as SSH host)")
    parser.add_argument('--disable-ssh-vsock', action='store_true',
                        help="Disable connections to SSH using AF_VSOCK")
    parser.add_argument('domain_name', help="Libvirt domain name of VM")
    parser.add_argument('app', nargs=argparse.REMAINDER,
                        help="App to launch")
    args = parser.parse_args()

    system_hostname = socket.gethostname()
    domain_name = args.domain_name

    if args.skip_libvirt:
        xpra_host = domain_name
        user = None
        identity_file = None
    else:
        split_dname = domain_name.split('@', 1)
        if len(split_dname) > 1:
            user = split_dname[0]
            domain_name = split_dname[1]
        else:
            user = args.user

        if not args.disable_ssh_vsock:
            xpra_host = domain_name
        else:
            xpra_host = domain_to_addr(domain_name)
        identity_file = args.identity

    if args.cid is not None:
        cid = args.cid
    else:
        cid = domain_to_cid(domain_name)

    xpra_display = 10
    xpra_title = "@title@ on {0}".format(domain_name)

    ssh_args = ['ssh']
    scp_args = ['scp']
    if identity_file is not None:
        ssh_args += ['-i', identity_file]
        scp_args += ['-i', identity_file]
    if user is not None:
        ssh_args += ['-l', user]
    if not args.skip_libvirt:
        hka_arg = [
            '-o',
            f'HostKeyAlias=appvm:{system_hostname}:{domain_name}'
        ]
        ssh_args += hka_arg
        scp_args += hka_arg
    if not args.disable_ssh_vsock:
        pc_arg = [
            '-o',
            f'ProxyCommand=socat vsock-connect:{cid}:22 -',
        ]
        ssh_args += pc_arg
        scp_args += pc_arg

    if args.shell:
        subprocess.call(ssh_args + [xpra_host] + args.app)
    else:
        ssh_cmd = ' '.join(ssh_args)

        if len(args.app) <= 0:
            app = "firefox"
        else:
            app = args.app[0]

        p = xpra_attach(cid, xpra_title)
        if p != 0:
            scp_if_exists('.config/gtk-2.0/gtkrc', xpra_host, user, scp_args)
            scp_if_exists('.config/gtk-3.0/settings.ini', xpra_host, user,
                          scp_args)

            subprocess.call(ssh_args + [xpra_host] + [
                'xpra', 'start', ':{0}'.format(xpra_display),
                '--bind-vsock=auto:14500',
                '--start-after-connect={0}'.format(shlex.quote(app)),
            ])
            time.sleep(5)
            xpra_attach(cid, xpra_title)
