#!/usr/bin/python3
#
# This script is run by a VM host (eg "west") to prepare itself for testing

import os
import sys
import socket
import shutil
import subprocess
import pexpect
import glob
from pathlib import Path
import re
import argparse
from enum import Enum

def path_touch(dst, mode=None):
    if mode: #it is seems unhappy with mode=None
        Path(dst).touch(mode)
    else:
        Path(dst).touch()

def rsync_ap(src, dst, timer=20):

    cmd = "/usr/bin/rsync --delete -q -aP"
    cmd += " %s %s" % (src, dst)

    try:
        output = subprocess.check_output(cmd, shell=True, timeout=timer, stderr=subprocess.STDOUT)
    except subprocess.TimeoutExpired:
        print( "EXCEPTION TIMEOUT ? cmd %s , cwd %s" % (os.getcwd(), cmd))
    except subprocess.CalledProcessError as e:
         print ("EXCEPTION ? cwd %s , %s %s" % (os.getcwd(), cmd, e.output))

def mount_bind(src, dst, mode=None, touch_src_file=False, mkdir_src=False, wipe_old_dst=False, fatal=True):

    if touch_src_file and mkdir_src:
        print("conflicting options touch_src_file and mkdir_src");

    if mkdir_src and not os.path.isdir(src) and not touch_src_file:
        os.makedirs(src)

    if touch_src_file :
        path_touch(src, mode)

    if wipe_old_dst:
        wipe_old(dst)

    if os.path.isdir(src):
        if not os.path.isdir(dst) and not os.path.islink(dst):
            os.makedirs(dst)
    elif os.path.isfile(src):
        if not os.path.isfile(dst) and not os.path.islink(dst):
            path_touch(dst, mode)
    else:
        mode_str = ''
        if mode:
            mode_str = "mode 0o%o" % mode
        print("mount_bind unknown action src=%s dst=%s %s"
                 % (src, dst, mode_str))
        if fatal:
            sys.exit(1)
        return True

    cmd = ['/bin/mount', '--bind',  src, dst]
    print(f"namespace: {cmd}")
    o = subprocess.run(cmd, stderr=subprocess.PIPE, stdout=subprocess.PIPE,
                       encoding='utf-8')
    if o.returncode:
        print("mount_bind failed %s" % cmd)
        if fatal:
            sys.exit(1)

    return o.returncode

def lsw_init_dir(src, dst):
    if args.namespace:
        print(f"namespace: bind {src} {dst}")
        return mount_bind(src, dst, mkdir_src=True,  wipe_old_dst=True, fatal=True)
    else:
        os.makedirs(dst, exist_ok=True)
    return

def lsw_cp_file(src, dst, mode=None, nsbasepath=''):

    if args.namespace:
        ns_dst="%s/%s" %(nsbasepath, dst) #copy to local NS/hostname/<path> then mount
        print(f"namespace: copy {src} {ns_dst}")
        shutil.copy(src, ns_dst)
        if mode:
            print(f"namespace: chmod {mode} {ns_dst}")
            os.chmod(ns_dst, mode)
        print(f"namespace: bind {ns_dst} {dst}")
        mount_bind(ns_dst, dst)
    else:
        shutil.copy(src, dst)
        if mode:
            os.chmod(dst, mode)
    return

def umount(dst):
   if os.path.islink(dst):
       print(f"namespace: unlink {dst}")
       os.unlink(dst)
   else:
       cmd = "umount %s" % dst
       print(f"namespace: {cmd}")
       subprocess.run(cmd, shell=True, capture_output=True, encoding="ascii",
                      check=True)

def umount_all(dst):
    o = subprocess.run("mount", shell=True, capture_output=True, encoding="ascii", check=True)
    for line in o.stdout.splitlines():
        mdst = line.split(' ')[2]
        if mdst == dst:
            umount(dst)

def wipe_old(dst):
    if args.namespace:
        print(f"namespace: wipe_old {dst}")
        umount_all(dst)
    elif os.path.islink(dst) or os.path.isfile(dst):
        os.unlink(dst)
    elif os.path.isdir(dst):
        shutil.rmtree(dst)
    # else: #no existing file
    #   print("cannot wipe unknown type %s" % dst)

def shell(command, out=False):
    """Run command as a sub-process and report failures"""
    if args.verbose:
        print(command)
    status, output = subprocess.getstatusoutput(command)
    if status:
        print(f"command '{command}' failed with status {status}: {output}")
    elif (args.verbose or out) and output:
        print(output)
    return output

def choose_config_file(hostname, testpath, config_path, strict, default=None):
    (config_absdir, config_file) = os.path.split(config_path)
    if os.path.isabs(config_absdir):
        config_dir = config_absdir[1:] # drop slash; better?
    else:
        config_dir = config_absdir
    # extension contains the dot
    (config_base, extension) = os.path.splitext(config_file)

    # pass 1: look in the test directory, reject ambiguous files.
    # When exact, don't match HOSTNAME.EXTENSION.
    paths = []
    paths.append((False, os.path.join(testpath, hostname + "." + config_file)))
    paths.append((True, os.path.join(testpath, hostname + extension)))
    paths.append((None, os.path.join(testpath, config_file)))
    path = None
    for fuzzy, p in paths:
        if os.path.isfile(p):
            if strict is True and fuzzy is True:
                # strict so ignore fuzzy match
                continue
            if path:
                sys.exit("conflicting files %s %s in test directory" % (p, path))
            path = p
    if path:
        return path

    # pass 2: look in default
    if default:
        if not os.path.isfile(default):
            sys.exit(f"the default {default} for {config_path} is not a file")
        return default

    # dump all the paths?
    return None

def get_configsetup(key):
    # without --config /dev/null most tests would pass.  However, the
    # test that follows a test with broken config will fail And during
    # a testrun the worker that hit error may not recover.
    o = subprocess.run(["ipsec", "addconn", f"--configsetup={key}", "--config", "/dev/null"],
                       capture_output=True, encoding="ascii", check=True,)
    return o.stdout.strip()

def copy_config_file(hostname, testpath, config_path,
                     strict=False, optional=False,
                     nsprefix="", default=None):
    src = choose_config_file(hostname, testpath, config_path, strict,
                             default=default)
    if src:
        # config_absdir has a leading / so below is always valid
        dst = nsprefix + config_path
        dstdir = os.path.dirname(dst)
        if not os.path.isdir(dstdir):
            os.mkdir(dstdir)
        # When NS DST=<host>/NS/<test>/CONFIG_PATH
        shutil.copy(src, dst)
        # Now mount DST=<host>/NS/<test>/CONFIG_PATH over CONFIGI_PATH
        if nsprefix: # non-None and non-Empty
            mount_bind(dst, config_path);
    elif not optional:
        print("ERROR: %s's %s not found" % (hostname, config_path))

# prevent double installs from going unnoticed
if os.path.isfile("/usr/libexec/ipsec/pluto"):
    if os.path.isfile("/usr/local/libexec/ipsec/pluto"):
        sys.exit("\n\n---------------------------------------------------------------------\n"
                 "ABORT: found a swan userland in the base system as well as /usr/local\n"
                 "---------------------------------------------------------------------\n")

parser = argparse.ArgumentParser(description='swan-prep arguments')

parser.add_argument('--testpath', '-t', action='store',
                    default=os.getcwd(), help="Test directory full path %s " % os.getcwd())
parser.add_argument('--hostname', '-H', action='store',
                    default='', help='The name of the host to prepare as')
# we should get this from the testparams.sh file?
parser.add_argument('--userland', '-u', action='store',
                    default='libreswan', help='which userland to prepare')
parser.add_argument('--strongswan-version', action='store',
                    default='', help='strongswan version expect')

class Keys(Enum):
    X509 = 'x509'
    HOST = 'host'
    NONE = 'none'
parser.add_argument('--keys',
                    choices=[Keys.X509, Keys.HOST, Keys.NONE], default=None,
                    help='create an NSS database containing X.509, host, or no keys')
parser.add_argument('--nokeys',
                    action='store_const', dest='keys', const=Keys.NONE,
                    help='create an empty NSS database')
parser.add_argument('--hostkeys',
                    action='store_const', dest='keys', const=Keys.HOST,
                    help="create an NSS database containing the host keys")
parser.add_argument('--x509keys', '-x', '--x509',
                    action='store_const', dest='keys', const=Keys.X509,
                    help='create an NSS database containing X.509 certificates and private keys')

nsspassword="s3cret"
parser.add_argument('--nsspw', action='store_true',
                    help='set the security password on the NSS database to '+nsspassword)

parser.add_argument('--46', '--64', action='store_true',
                    help='Do not disable IPv6. Default is disable IPv6 ', dest='ipv46', default=False)
parser.add_argument('--verbose', '-v', action='store_true',
                    help='more verbose')
parser.add_argument('--namespace', action='store_true',
        default='', help='Running inside name sapace')

args = parser.parse_args()

if args.hostname:
    hostname = args.hostname
else:
    hostname = socket.gethostname()
if "." in hostname:
    hostname = hostname.split(".")[0]
if hostname == "nic":
    sys.exit("swan-prep is never run on NIC")

# Validate this is sane?
testpath = args.testpath
if not os.path.isdir(testpath):
    sys.exit("Unknown or bad testpath '%s'" % args.testname)

testname = os.path.basename(testpath)

# nsrun set magic string SWAN_PLUTOTEST=YES in the environment
# to identify namespace
# echo $SWAN_PLUTOTEST
if args.namespace:
    print(f"namespace: --namespace")
elif "SWAN_PLUTOTEST" in os.environ:
    print(f"namespace: SWAN_PLUTOTEST")
    args.namespace = True
else:
    o = subprocess.run('ip netns identify', shell=True, capture_output=True, encoding="ascii", check=True,)
    namespace = o.stdout
    namespace_expect = hostname + '-' + testname
    if re.search(hostname, namespace):
        print(f"namespace: ip netns identity")
        args.namespace = True

resolve_conf = "/etc/resolv.conf"
etc_hosts = "/etc/hosts"

ipsecdir = get_configsetup("ipsecdir")
configfile = get_configsetup("configfile")
secretsfile = get_configsetup("secretsfile")
exampledir = get_configsetup("exampledir")

if args.namespace:
    nsbasepath ="%s/NS/%s" % (testpath, hostname) #will create testpath/NS/hostname/*
    if not os.path.isdir(nsbasepath):
        print(f"namespace: mkdir {nsbasepath}")
        os.makedirs(nsbasepath)
else:
    nsbasepath = ''

# Setup pluto.log softlink and bindmount in namespace

subprocess.run(["ipsec", "stop"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, encoding="ascii") # gentle try before wiping files and killall

outputdir = "%s/OUTPUT/" % testpath
if not os.path.isdir(outputdir):
    os.mkdir(outputdir)
    os.chmod(outputdir, 0o777)

match args.userland:
    case "libreswan":
        dname = "pluto"
    case "strongswan":
        dname = "charon"

daemonlogfile = "%s/%s.%s.log" % (outputdir, hostname, dname)
tmplink = "/tmp/%s.log" % dname
wipe_old(tmplink)

Path(daemonlogfile).touch(mode=0o777)

if args.namespace :
    print(f"namespace: {daemonlogfile} {tmplink}")
    mount_bind(daemonlogfile, tmplink, touch_src_file=True)
else :
    os.symlink(daemonlogfile, tmplink)

match args.userland:
    case "libreswan" | "strongswan":
        None
    case _:
        sys.exit("swan-prep: unknown userland type '%s'" % args.userland)

# print "swan-prep running on %s for test %s with userland
# %s"%(hostname,testname,userland)

# wipe or unmount any old configs
# do not clean pid files or files opened by pluto, charon .. yet.
wipe_old(configfile)
wipe_old(secretsfile)
shutil.rmtree(f"{ipsecdir}/policies", ignore_errors=True);
wipe_old("/etc/strongswan")

if args.namespace:
    wipe_old(resolve_conf)
    wipe_old(etc_hosts)

# if using systemd, ensure we don't restart pluto on crash
if os.path.isfile("/lib/systemd/system/ipsec.service") and not args.namespace:
    service = "".join(open("/lib/systemd/system/ipsec.service").readlines())
    if "Restart=always" in service:
        fp = open("/lib/systemd/system/ipsec.service", "w")
        fp.write("".join(service).replace("Restart=always", "Restart=no"))
    # always reload to avoid "service is masked" errors
    subprocess.getoutput("/bin/systemctl daemon-reload")

# AA_201902 what happens to audit log namespace
# we have to cleanup the audit log or we could get entries from previous test
if os.path.isfile("/var/log/audit/audit.log"):
    fp = open("/var/log/audit/audit.log", "w")
    fp.close()
    if os.path.isfile("/lib/systemd/system/auditd.service"):
        subprocess.getoutput("/bin/systemctl restart auditd.service")

# disable some daemons that could cause noise packets getting encrypted
if not args.namespace :
    subprocess.getoutput("/bin/systemctl stop chronyd")
    subprocess.getoutput("/bin/systemctl stop sssd")

# work around for ACQUIRE losing sub-type, causing false positive test results
subprocess.getoutput('sysctl net.ipv4.ping_group_range="1 0"')

# ensure cores are just dropped, not sent to aobrt-hook-ccpp or systemd.
# setting /proc/sys/kernel/core_pattern to "core" or a pattern does not
# work. And you cannot use shell redirection ">". So we hack it with "tee"
pattern = "|/usr/bin/tee /tmp/core.%h.%e.%p"
fp = open("/proc/sys/kernel/core_pattern", "w")
fp.write(pattern)
fp.close()

match args.userland:
    case "strongswan":
        dst = "/etc/strongswan"
    case _:
        dst = ipsecdir

src = "%s%s" % (nsbasepath, dst)
lsw_init_dir(src, dst)

match args.userland:
    case "libreswan":
        dst = "/run/pluto"
    case "strongswan":
        dst="/run/strongswan"
    case _:
        sys.exit("Unknown run directory")

src = "%s%s" % (nsbasepath, dst)
lsw_init_dir(src, dst)

match args.userland:

    case "libreswan":

        # fill in any missing dirs
        os.mkdir("/etc/ipsec.d/policies")

        copy_config_file(hostname, testpath, "/etc/ipsec.conf",
                         default=os.path.join(exampledir, "ipsec.conf-sample"),
                         nsprefix=nsbasepath)
        copy_config_file(hostname, testpath, "/etc/ipsec.secrets",
                         default=os.path.join(exampledir, "ipsec.secrets-sample"),
                         nsprefix=nsbasepath)

    case "strongswan":

        # check version and spew all over test output
        output = subprocess.getoutput("strongswan version")
        strongswan_version = "U6.0.7"
        if args.strongswan_version:
            strongswan_version = args.strongswan_version
        if not strongswan_version in output:
            print("strongswan %s must be installed" % (strongswan_version))
            print("")
            print(output)

        # required to write log file in /tmp
        subprocess.getoutput("setenforce 0")

        for directory in ("ipsec.d/",
                          "ipsec.d/aacerts/",
                          "ipsec.d/ocspcerts/",
                          "ipsec.d/cacerts/",
                          "ipsec.d/private/",
                          "ipsec.d/certs/",
                          "swanctl/",
                          "swanctl/x509",
                          "swanctl/x509ac",
                          "swanctl/x509ca",
                          "swanctl/x509ocsp",
                          "swanctl/x509aa",
                          "swanctl/x509crl",
                          "swanctl/pubkey",
                          "swanctl/private",
                          "swanctl/rsa",
                          "swanctl/ecdsa",
                          "swanctl/pkcs8",
                          "swanctl/pkcs12"):
            dst = os.path.join("/etc/strongswan", directory)
            if os.path.isdir(dst):
                shutil.rmtree(dst)
            os.mkdir(dst)
        copy_config_file(hostname, testpath, "/etc/strongswan/strongswan.conf", strict=True)
        copy_config_file(hostname, testpath, "/etc/strongswan/ipsec.conf", optional=True)
        copy_config_file(hostname, testpath, "/etc/strongswan/ipsec.secrets", optional=True)
        copy_config_file(hostname, testpath, "/etc/strongswan/swanctl/swanctl.conf", strict=True, optional=True)

# test specific files
xauthpasswd = "%s/%s.passwd" % (testpath, hostname)

if os.path.isfile(xauthpasswd):
    lsw_cp_file(xauthpasswd, "/etc/ipsec.d/passwd", nsbasepath=nsbasepath)

# restore /etc/hosts to original - some tests make changes
lsw_cp_file("/testing/baseconfigs/all/etc/hosts", "/etc/hosts", nsbasepath=nsbasepath)
resolv = "/testing/baseconfigs/all/etc/resolv.conf"
if os.path.isfile("/testing/baseconfigs/%s/etc/resolv.conf" % hostname):
    resolv = "/testing/baseconfigs/%s/etc/resolv.conf" % hostname
else:
    resolv = "/testing/baseconfigs/all/etc/resolv.conf"
dst = resolve_conf
if not args.namespace and os.path.islink(dst): # on fedora 22 it is link first remove the link
    os.unlink(dst)
lsw_cp_file(resolv, "/etc/resolv.conf", nsbasepath=nsbasepath)

sysconfigd = "%s/etc/sysconfig" % (nsbasepath)
if not os.path.isdir(sysconfigd):
    os.makedirs(sysconfigd)

if args.userland in ("libreswan") and not args.namespace:
    lsw_cp_file("/testing/baseconfigs/all/etc/sysconfig/pluto", "/etc/sysconfig/pluto",
                nsbasepath=nsbasepath)

# Set up NSS DB

if args.userland in ("libreswan"):

    # clean up existing
    nssdir = "/var/lib/ipsec/nss"
    shell(f"rm -rf {nssdir}/*")

    match args.keys:
        case Keys.HOST:
            # This brings in the pre-generated nss *.db files that contain
            # raw host keys.  Basic testing tends to use pre-shared keys
            # and not these files.
            print("Creating NSS database containing host keys")
            for dbfile in ("cert9.db", "key4.db"):
                src = os.path.join("/testing/baseconfigs", hostname, "etc/ipsec.d", dbfile)
                dst = os.path.join(nssdir, dbfile)
                shutil.copyfile(src, dst)
                if not args.namespace:
                    os.chown(dst, 0, 0)

        case Keys.X509:
            if not os.path.isfile("/testing/x509/pki/real/mainca/root.p12"):
                print("\n\n---------------------------------------------------------------------\n"
                      "WARNING: no root.p12 file, did you run './kvm keys'?\n"
                      "---------------------------------------------------------------------\n")
            shell(f"ipsec initnss", out=True)
            shell(f"/testing/x509/import.sh real/mainca/{hostname}.p12", out=True)

        case Keys.NONE:
            shell(f"ipsec initnss", out=True)

    # Determine the NSS password and save; assumes above set up the
    # NSS database.  Assumes machine isn't yet in FIPS mode?

    if args.nsspw:
        with open("/run/pluto/nsspw", "w") as f:
            f.write(nsspassword)
            f.write("\n")
        # Store the password in "$NSSDIR/nsspassword" so that pluto
        # can use it to open the the NSS DB
        with open(ipsecdir + "/nsspassword", "w") as f:
            if args.nsspw:
                f.write(f"NSS Certificate DB:{nsspassword}\n")
        with open("/tmp/pw", "w") as f:
            f.write("\n")
        shell(f"/usr/bin/certutil -W -f /tmp/pw -@ /run/pluto/nsspw -d sql:{nssdir}", out=True)

if args.namespace:
    # good idea for namespace. KVM sysctl.conf get copid.
    subprocess.getoutput("sysctl -p /testing/baseconfigs/all/etc/sysctl.conf")

if not args.ipv46:
    subprocess.getoutput("sysctl net.ipv6.conf.all.disable_ipv6=1")
    subprocess.getoutput("sysctl net.ipv6.conf.default.disable_ipv6=1")

subprocess.getoutput("iptables -F");
subprocess.getoutput("iptables -X");

if not args.namespace: # inside namespace this would kill pluto from other ns
    # shouldn't happen early on? now wiped run time files pid etc.
    # this is probably a last resort? may be a more gentle attempt at the beginning.
    #
    # final prep - this kills any running userland
    subprocess.call(["systemctl", "stop", "ipsec"])
    # for some reason this fails to stop strongswan?
    subprocess.call(["systemctl", "stop", "strongswan"])
    # python has no pidof - just outsource to the shell, thanks python!
    for dname in ( "pluto", "charon", "starter", "iked" ):
        try:
            if args.verbose:
                print ("INFO found daemon running stop it %s" % dname)
            subprocess.check_output(["killall", "-9", dname], stderr=subprocess.STDOUT)
        except:
            pass

    if os.path.isfile("/usr/sbin/getenforce"):
        selinux = subprocess.getoutput("/usr/sbin/getenforce")
        if os.path.isfile("/usr/sbin/restorecon") and selinux == 'Enforcing':
            subprocess.getoutput("restorecon -Rv /etc/ipsec.* /var/lib/ipsec /usr/local/libexec/ipsec /usr/local/sbin/ipsec")

for pidfile in ("/run/pluto/pluto.pid", "/run/strongswan/charon.pid", "/run/spmd.pid", ):
    if os.path.isfile(pidfile):
        os.unlink(pidfile)
