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

# Name of the socket used for connecting to libvirtd
# https://gitlab.com/libvirt/libvirt/-/blob/65312001bd972df8b7d4f11ea4662aff4889bee5/src/remote/remote_sockets.c#L144-147
LIBVIRT_SOCK = '/var/run/libvirt/libvirt-sock'

# Regex matching the ssh tunnel script generated by modern versions
# of libvirt, which prefer using virt-ssh-helper instead of netcat
# https://gitlab.com/libvirt/libvirt/-/blob/65312001bd972df8b7d4f11ea4662aff4889bee5/src/rpc/virnetclient.c#L446-448
LIVE_MIGRATION_TUNNEL_HELPER = "virt-ssh-helper +(qemu:///system|'qemu:///system')"

# Regex matching the ssh tunnel script generated by old versions of
# libvirt or when proxy=netcat is part of the connection URI
# https://gitlab.com/libvirt/libvirt/-/blob/65312001bd972df8b7d4f11ea4662aff4889bee5/src/rpc/virnetclient.c#L437-444
LIVE_MIGRATION_TUNNEL_NETCAT = f"(nc|'nc') +\$ARG +-U +{LIBVIRT_SOCK}"

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 re.search(LIVE_MIGRATION_TUNNEL_HELPER, command):
    args = ['virt-ssh-helper', 'qemu:///system']
    allow_command('nova', args)
if re.search(LIVE_MIGRATION_TUNNEL_NETCAT, command):
    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)
