#!/usr/libexec/platform-python
import os
import sys
import syslog

# Handle libvirt ssh tunnel script snippet
# https://github.com/libvirt/libvirt/blob/f0803dae93d62a4b8a2f67f4873c290a76d978b3/src/rpc/virnetsocket.c#L890
LIBVIRT_SOCK = '/var/run/libvirt/libvirt-sock'
# Need to keep this in case proxy=netcat is part of the uri string
# https://github.com/libvirt/libvirt/blob/108676c225c8aeb49bbbd5b8e55f7dbfedc71ac0/src/rpc/virnetclient.c#L437-L444
LIVE_MIGRATION_TUNNEL_NETCAT = (
    "if 'nc' -q 2>&1 | grep \"requires an argument\" >/dev/null 2>&1; then "
    "ARG=-q0;"
    "else "
    "ARG=;"
    "fi;"
    f"'nc' $ARG -U {LIBVIRT_SOCK}")

# The virt-ssh-helper command string includes the origin netcat one (without the sh)
# https://github.com/libvirt/libvirt/blob/108676c225c8aeb49bbbd5b8e55f7dbfedc71ac0/src/rpc/virnetclient.c#L452-L457
LIVE_MIGRATION_VIRT_SSH_HELPER = (
    "sh -c 'which virt-ssh-helper 1>/dev/null 2>&1; if test $? = 0; then     " \
    "virt-ssh-helper 'qemu:///system'; " \
    f"else    {LIVE_MIGRATION_TUNNEL_NETCAT}; fi'")
# We need to put back the sh in the netcat command
LIVE_MIGRATION_TUNNEL_NETCAT_SH = f"sh -c '{LIVE_MIGRATION_TUNNEL_NETCAT}'"

COLD_MIGRATION_ROOT = '/var/lib/nova/instances/'
COLD_MIGRATION_CMDS = [
    ['mkdir', '-p'],
    ['rm', '-rf'],
    ['touch'],
    ['rm'],
    ['scp', '-r', '-t'],
    ['scp', '-r', '-f'],
    ['scp', '-t'],
    ['scp', '-f'],
    ['/usr/libexec/openssh/sftp-server'],
]
ROOTWRAP_ARGS = ['/usr/bin/nova-rootwrap', '/etc/nova/migration/rootwrap.conf']

command = os.environ.get('SSH_ORIGINAL_COMMAND')
ssh_connection = os.environ.get('SSH_CONNECTION')
if command is None:
    sys.stderr.write('This command must be run via SSH ForceCommand (see man 5 sshd_config).\n')
    sys.exit(1)

syslog.openlog('nova_migration_wrapper')

def allow_command(user, args):
    syslog.syslog(syslog.LOG_INFO, "Allowing connection='{}' command={} ".format(
        ssh_connection,
        repr(args)
    ))
    os.execlp('sudo', 'sudo', '-u', user, *args)

def deny_command(args):
    syslog.syslog(syslog.LOG_ERR, "Denying connection='{}' command={}".format(
        ssh_connection,
        repr(args)
    ))
    sys.stderr.write('Forbidden\n')
    sys.exit(1)

def validate_cold_migration_cmd(args):
    target_path = os.path.normpath(args[-1])
    # We simply remove the last argument to get the command with various args
    cmd = args[:-1]
    cmd_with_args = cmd in COLD_MIGRATION_CMDS and target_path.startswith(COLD_MIGRATION_ROOT)
    # In the case of a command with no args, cmd will be empty
    cmd_no_args = not cmd and args in COLD_MIGRATION_CMDS
    return cmd_with_args or cmd_no_args

# TODO(dvd): Move this in TripleO
# Rules
args = command.split(' ')
if command == LIVE_MIGRATION_VIRT_SSH_HELPER:
    args = ['virt-ssh-helper', 'qemu:///system']
    allow_command('nova', args)
if command == LIVE_MIGRATION_TUNNEL_NETCAT_SH:
    args = ['nc', '-U', LIBVIRT_SOCK]
    allow_command('nova', args)
if validate_cold_migration_cmd(args):
    args = ROOTWRAP_ARGS + args
    allow_command('root', args)
deny_command(args)
