#!/usr/bin/env python3
#
# Copyright VyOS maintainers and contributors <maintainers@vyos.io>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 or later as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# File: check-oci-container
# Purpose:
#  This script verifies the OCI container image produced by scripts/iso-to-oci.
#  It imports the rootfs tarball into the local container runtime, boots it with
#  systemd as PID 1, waits for vyos-router to come up and then attaches to the
#  VyOS CLI to run a set of basic operational and configuration mode tests.
#
#  It also asserts the rootfs modifications iso-to-oci performs (locale, login
#  banner, /config symlink, removal of container incompatible CLI nodes) are
#  still in place - those are invisible until someone imports the image by hand.
#
#  Arguments:
#    [--image]      OCI tarball to test, defaults to the newest one in $PWD
#    [--runtime]    Force 'docker' or 'podman', default is to auto-detect
#    [--name]       Container name, autogenerated if not given
#    [--keep]       Keep container and imported image after completion
#    [--logfile]    name of logfile to save, defaulting to stdout
#    [--silent]     only print on errors
#    [--debug]      print all communication with the container
#
#  Exit codes: 0 on success, 1 if the image failed a test, 2 if the test could
#  not be run at all (no runtime, no image) - same convention as iso-to-oci.

import sys
import os
import time
import argparse
import subprocess
import random
import traceback
import logging
import re
import shutil
import platform

from datetime import datetime
from glob import glob

import pexpect

EXCEPTION = 0

def bail(message: str):
    """ Preconditions for running the test at all are not met - this is not a
        test failure, so exit 2 like scripts/iso-to-oci does """
    print(f'E: {message}', file=sys.stderr)
    sys.exit(2)

# The container is always started with --hostname vyos so the very same prompt
# regexes used by check-qemu-install remain valid
op_mode_prompt = r'vyos@vyos:~\$'
cfg_mode_prompt = r'vyos@vyos#'
default_user = 'vyos'

# Time (in seconds) granted to systemd inside the container to reach a settled
# state and bring up vyos-router
boot_timeout = 300

# containerlab uses the image healthcheck to determine when a VyOS node is
# ready - keep in sync with the import hint printed by scripts/iso-to-oci
healthcheck = 'HEALTHCHECK --start-period=10s CMD systemctl is-system-running'

# Tarballs are named vyos-<version>-oci-<arch>.tar.xz by scripts/iso-to-oci
IMAGE_RE = re.compile(r'^vyos-(?P<version>.+)-oci-(?P<arch>[^-]+)\.tar\.xz$')

# Map the host architecture to the Debian architecture used in the image name
ARCH_MAP = {
    'x86_64': 'amd64',
    'amd64': 'amd64',
    'aarch64': 'arm64',
    'arm64': 'arm64',
}

parser = argparse.ArgumentParser()
parser.add_argument('--image', help='OCI rootfs tarball to test')
parser.add_argument('--runtime', help='Container runtime to use',
                    choices=['docker', 'podman'])
parser.add_argument('--name', help='Name of the test container')
parser.add_argument('--keep', help='Do not remove container and image after testing',
                    action='store_true', default=False)
parser.add_argument('--silent', help='Do not show output on stdout unless an error has occurred',
                    action='store_true', default=False)
parser.add_argument('--debug', help='Send all debug output to stdout',
                    action='store_true', default=False)
parser.add_argument('--logfile', help='Log to file')
args = parser.parse_args()

host_arch = ARCH_MAP.get(platform.machine())
if host_arch is None:
    bail(f'Unsupported host architecture "{platform.machine()}" - this script '
         'only runs on amd64 or arm64 hosts.')

class StreamToLogger(object):
    """ Fake file-like stream object that redirects writes to a logger instance. """
    def __init__(self, logger, log_level=logging.INFO):
        self.logger = logger
        self.log_level = log_level
        self.linebuf = ''
        self.ansi_escape = re.compile(r'\x1B\[[0-?]*[ -/]*[@-~]')

    def write(self, buf):
        temp_linebuf = self.linebuf + buf
        self.linebuf = ''
        for line in temp_linebuf.splitlines(True):
            if line[-1] == '\n':
                tmp = self.ansi_escape.sub('', line.rstrip())
                if tmp:
                    self.logger.log(self.log_level, tmp)
            else:
                self.linebuf += line

    def flush(self):
        if self.linebuf != '':
            tmp = self.ansi_escape.sub('', self.linebuf.rstrip())
            if tmp:
                self.logger.log(self.log_level, tmp)
        self.linebuf = ''

# Setting up logger
log = logging.getLogger()
log.setLevel(logging.DEBUG)

stl = StreamToLogger(log, logging.DEBUG)
formatter = logging.Formatter('%(levelname)5s - %(message)s')

handler = logging.StreamHandler(sys.stdout)
if args.silent:
    handler.setLevel(logging.ERROR)
elif args.debug:
    handler.setLevel(logging.DEBUG)
else:
    handler.setLevel(logging.INFO)

handler.setFormatter(formatter)
log.addHandler(handler)

if args.logfile:
    file_handler = logging.FileHandler(filename=args.logfile, mode='w')
    file_handler.setLevel(logging.DEBUG)
    file_handler.setFormatter(formatter)
    log.addHandler(file_handler)

def find_image() -> str:
    """ Return the newest OCI tarball matching the host architecture """
    images = glob(f'vyos-*-oci-{host_arch}.tar.xz')
    if not images:
        bail(f'No OCI image found - run "make oci" first, or pass --image. '
             f'Looked for vyos-*-oci-{host_arch}.tar.xz in {os.getcwd()}')
    return max(images, key=os.path.getmtime)

def find_runtime() -> str:
    """ Prefer docker, fall back to podman """
    if args.runtime:
        if not shutil.which(args.runtime):
            bail(f'Requested container runtime "{args.runtime}" not found in PATH.')
        return args.runtime
    for runtime in ['docker', 'podman']:
        if shutil.which(runtime):
            return runtime
    bail('Neither docker nor podman found in PATH - a container runtime is '
         'required to test the OCI image. Note the vyos-build container itself '
         'ships no runtime, run this on the host.')

# Importing a compressed rootfs tarball is slow - the runtime decompresses the
# ~400 MiB xz stream more than once, which takes well over ten minutes
import_timeout = 3600

def run(cmd: list, check=True, timeout=120) -> subprocess.CompletedProcess:
    """ Run a container runtime command and log its output """
    log.debug(f'Executing command: {" ".join(cmd)}')
    ret = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
    if ret.stdout.strip():
        log.debug(ret.stdout.strip())
    if ret.stderr.strip():
        log.debug(ret.stderr.strip())
    if check and ret.returncode != 0:
        raise Exception(f'Command "{" ".join(cmd)}" failed with exit code '
                        f'{ret.returncode}: {ret.stderr.strip()}')
    return ret

def container_exec(runtime: str, name: str, command: str, check=True) -> str:
    """ Non-interactive command execution inside the container """
    ret = run([runtime, 'exec', name, '/bin/bash', '-c', command], check=check)
    return ret.stdout.strip()

def get_import_cmd(runtime: str, tag: str) -> list:
    """ docker takes a combined --platform=<os>/<arch>, podman only knows the
        separate --os/--arch flags and errors out on --platform """
    cmd = [runtime, 'import']
    if runtime == 'podman':
        cmd += ['--os', 'linux', '--arch', image_arch]
    else:
        cmd += [f'--platform=linux/{image_arch}']
    cmd += [image, tag, '--change', 'CMD ["/sbin/init"]',
            '--change', healthcheck]
    return cmd

def verify_healthcheck(runtime: str, tag: str) -> None:
    """ The healthcheck is not part of the rootfs tarball, it is attached by
        "import --change" - a runtime silently dropping it would leave
        containerlab waiting for a node that never turns healthy. Inspect the
        raw JSON as docker and podman expose the healthcheck under different
        keys """
    ret = run([runtime, 'image', 'inspect', tag])
    if 'is-system-running' not in ret.stdout:
        raise Exception(f'{runtime} did not store the image healthcheck '
                        f'"{healthcheck}" - inspect output: {ret.stdout}')

def get_run_cmd(runtime: str, name: str, tag: str) -> list:
    """ systemd as PID 1 needs a writable /run and cgroup access - podman knows
        how to do that on its own, docker has to be told explicitly """
    cmd = [runtime, 'run', '--detach', '--name', name, '--hostname', 'vyos',
           '--privileged']
    if runtime == 'podman':
        cmd += ['--systemd=always']
    else:
        cmd += ['--cgroupns=host',
                '--tmpfs', '/run',
                '--tmpfs', '/run/lock',
                '--volume', '/sys/fs/cgroup:/sys/fs/cgroup:rw',
                '--stop-signal', 'SIGRTMIN+3']
    cmd += [tag, '/sbin/init']
    return cmd

def wait_for_boot(runtime: str, name: str) -> None:
    """ Wait until systemd settled, vyos-router reports itself as active and the
        boot configuration has been applied. "degraded" is a perfectly valid end
        state in a container - services depending on kernel modules or firmware
        we stripped will fail. """
    log.info('Waiting for systemd to finish booting')
    deadline = time.time() + boot_timeout
    system_state = None
    while time.time() < deadline:
        # systemctl is-system-running exits non-zero for anything but "running"
        system_state = container_exec(runtime, name, 'systemctl is-system-running',
                                      check=False)
        if system_state in ['running', 'degraded']:
            break
        if system_state == 'maintenance':
            raise Exception('systemd entered maintenance mode inside the container')
        time.sleep(5)
    else:
        raise Exception(f'systemd did not finish booting within {boot_timeout} '
                        f'seconds, last state: "{system_state}"')

    log.info(f'systemd reached state "{system_state}"')
    if system_state == 'degraded':
        failed = container_exec(runtime, name,
                                'systemctl list-units --state=failed --no-legend --no-pager',
                                check=False)
        log.info(f'Failed units inside the container:\n{failed}')

    log.info('Waiting for vyos-router to become active')
    while True:
        if time.time() >= deadline:
            raise Exception('vyos-router did not become active inside the container')
        if container_exec(runtime, name, 'systemctl is-active vyos-router',
                          check=False) == 'active':
            log.info('vyos-router is active')
            break
        time.sleep(5)

    # vyos-router reports itself active before the boot configuration is fully
    # committed - the vyos user does not exist in the rootfs, it is created when
    # /usr/share/vyos/config.boot.default is applied. Wait for it explicitly,
    # otherwise a slow or broken commit surfaces as an opaque pexpect timeout.
    log.info(f'Waiting for user "{default_user}" to be created by the boot configuration')
    while True:
        if time.time() >= deadline:
            raise Exception(f'User "{default_user}" does not exist - vyos-router '
                            'did not apply the default configuration')
        uid = container_exec(runtime, name, f'id -u {default_user}', check=False)
        if uid:
            log.info(f'User "{default_user}" exists with uid {uid}')
            return
        time.sleep(5)

def expect_output(c, command: str, expected: str) -> None:
    """ Send a command and verify its output contains the expected string. The
        prompt is consumed afterwards so the session stays in sync. """
    c.sendline(command)
    c.expect(op_mode_prompt)
    output = c.before
    if expected not in output:
        raise Exception(f'Command "{command}" did not return "{expected}": {output.strip()}')

def basic_cli_tests(c, version: str) -> None:
    """ Container adapted version of the check-qemu-install testcase - no kernel
        version, GRUB or serial console checks, those do not apply to a container """
    log.info('Basic CLI configuration mode test')
    c.sendline('configure')
    c.expect(cfg_mode_prompt)
    c.sendline('exit')
    c.expect(op_mode_prompt)

    log.info('Basic CLI operational mode test')
    for command in ['show version', 'show version all | grep vyos-1x',
                    'show interfaces', 'show system memory',
                    'show system uptime', 'show configuration commands']:
        c.sendline(command)
        c.expect(op_mode_prompt)

    log.info(f'Verify running image reports version {version}')
    expect_output(c, 'jq -r .version /usr/share/vyos/version.json', version)

def iso_to_oci_tests(c) -> None:
    """ Verify the rootfs modifications performed by scripts/iso-to-oci """
    log.info('Verify locale is set to C.UTF-8')
    expect_output(c, 'echo "LANG=${LANG}"', 'LANG=C.UTF-8')

    log.info('Verify VyOS pre-login banner has been seeded')
    expect_output(c, 'cat /etc/issue', 'Welcome to VyOS')

    log.info('Verify /config symlink')
    expect_output(c, 'readlink /config', '/opt/vyatta/etc/config')

    log.info('Verify the podman runtime has been removed')
    expect_output(c, 'command -v podman netavark aardvark-dns || echo "no podman"',
                  'no podman')

    log.info('Verify container operational mode commands have been removed')
    expect_output(c, 'show container', 'Invalid command')

    log.info('Verify container incompatible CLI nodes have been removed')
    c.sendline('configure')
    c.expect(cfg_mode_prompt)
    for command in ['set container', 'set system console',
                    'set system option kernel']:
        c.sendline(command)
        c.expect(cfg_mode_prompt)
        # A removed configuration node is rejected with "Configuration path:
        # [...] is not valid", a removed CLI command with "Invalid command"
        output = c.before
        if 'is not valid' not in output and 'Invalid command' not in output:
            raise Exception(f'CLI node "{command}" still exists inside the '
                            f'container: {output.strip()}')
    c.sendline('exit')
    c.expect(op_mode_prompt)

image = args.image if args.image else find_image()
if not os.path.isfile(image):
    bail(f'OCI image file not found: {image}')

match = IMAGE_RE.match(os.path.basename(image))
if not match:
    bail(f'Unexpected OCI image filename "{os.path.basename(image)}" - expected '
         'vyos-<version>-oci-<arch>.tar.xz as generated by scripts/iso-to-oci')

vyos_version = match.group('version')
image_arch = match.group('arch')

runtime = find_runtime()
tag = f'vyos/vyos:{vyos_version}-oci-test'
name = args.name if args.name else \
    f'vyos-oci-test-{datetime.now().strftime("%Y%m%d%H%M%S")}-{random.randint(1000, 9999)}'

log.info(f'Testing OCI image {image} ({image_arch}) using {runtime}')

try:
    log.info(f'Importing OCI image as {tag} - this decompresses the whole '
             'rootfs and takes several minutes')
    run(get_import_cmd(runtime, tag), timeout=import_timeout)

    log.info('Verify image healthcheck used by containerlab is present')
    verify_healthcheck(runtime, tag)

    log.info(f'Starting container {name}')
    run(get_run_cmd(runtime, name, tag))

    wait_for_boot(runtime, name)

    log.info('Connecting to the VyOS CLI')
    c = pexpect.spawn(f'{runtime} exec -it {name} su - {default_user}',
                      logfile=stl, timeout=60,
                      encoding='utf-8', codec_errors='replace')
    c.expect(op_mode_prompt)
    log.info('Connected!')
    c.sendline('set terminal width 160')
    c.expect(op_mode_prompt)
    c.sendline('set terminal length 60')
    c.expect(op_mode_prompt)

    basic_cli_tests(c, vyos_version)
    iso_to_oci_tests(c)

    c.sendline('exit')
    c.close()

    log.info('OCI container image test finished successfully!')

except pexpect.exceptions.TIMEOUT:
    log.error('Timeout while waiting for the VyOS CLI')
    EXCEPTION = 1
except Exception as e:
    log.error(f'Unexpected error: {e}')
    log.error(traceback.format_exc())
    EXCEPTION = 1

finally:
    if EXCEPTION:
        log.info(f'Container logs of {name}:')
        run([runtime, 'logs', name], check=False)
    if args.keep:
        log.info(f'Keeping container {name} and image {tag} - connect using:')
        log.info(f'  {runtime} exec -it {name} su - {default_user}')
    else:
        log.info(f'Removing container {name} and image {tag}')
        run([runtime, 'rm', '--force', name], check=False)
        run([runtime, 'rmi', '--force', tag], check=False)

if EXCEPTION:
    log.error('Test failed')
    sys.exit(1)

sys.exit(0)