#!/usr/bin/env python3
"""Check rendered security relationships; optionally smoke-test that exact runtime."""
import base64, json, os, pathlib, socket, subprocess, sys, tempfile, time, urllib.request
import yaml
ROOT = pathlib.Path(__file__).resolve().parent
os.chdir(ROOT)
def run(*args, **kw):
    return subprocess.check_output(args, text=True, **kw).strip()
subprocess.run(['helm', 'lint', 'chart'], check=True)
docs = list(yaml.safe_load_all(run('helm', 'template', 'smoke', 'chart')))
objects = {x['kind']: x for x in docs}
pod = objects['Deployment']['spec']['template']['spec']
app = pod['containers'][0]
assert not pod['automountServiceAccountToken'] and not pod['enableServiceLinks']
assert pod['securityContext']['runAsNonRoot']
assert pod['securityContext']['seccompProfile']['type'] == 'RuntimeDefault'
assert pod['dnsPolicy'] == 'None' and not pod['dnsConfig'].get('searches')
assert all(not x.get(k) for k in ['hostNetwork', 'hostPID', 'hostIPC'] for x in [pod])
assert not any('hostPath' in x or 'persistentVolumeClaim' in x for x in pod['volumes'])
for container in pod['initContainers'] + pod['containers']:
    s = container['securityContext']
    assert s['readOnlyRootFilesystem'] and not s['allowPrivilegeEscalation']
    assert s['capabilities']['drop'] == ['ALL']
assert app['image'] == pod['initContainers'][0]['image']
assert objects['NetworkPolicy']['spec']['ingress'] == []
assert objects['NetworkPolicy']['spec']['egress'] == []
assert objects['Service']['spec']['type'] == 'ClusterIP'
assert objects['Service']['spec']['selector'] == objects['Deployment']['spec']['selector']['matchLabels']
assert set(['startupProbe', 'readinessProbe', 'livenessProbe']) <= app.keys()
assert not any(x['kind'] == 'Ingress' for x in docs)
print('Rendered security and service relationships passed', flush=True)
if '--runtime' not in sys.argv:
    sys.exit(0)
with tempfile.TemporaryDirectory(prefix='oldipmi-test-') as tmp:
    os.chmod(tmp, 0o755)
    dirs = {}
    for volume in pod['volumes']:
        path = pathlib.Path(tmp) / volume['name']
        path.mkdir(mode=0o777)
        path.chmod(0o777)
        dirs[volume['name']] = path
        if 'configMap' in volume:
            for name, data in objects['ConfigMap']['data'].items():
                (path / name).write_text(data)
    def container_args(container):
        args = ['podman', 'run', '--read-only', '--read-only-tmpfs=false', '--cap-drop=ALL',
                '--security-opt=no-new-privileges', '--userns=keep-id',
                '--user', str(container['securityContext'].get('runAsUser', pod['securityContext']['runAsUser']))]
        for mount in container['volumeMounts']:
            args += ['-v', str(dirs[mount['name']]) + ':' + mount['mountPath'] + (':ro' if mount.get('readOnly') else '')]
        return args
    init = pod['initContainers'][0]
    subprocess.run(container_args(init) + ['--rm', '--network=none', init['image']] + init['command'] + init['args'], check=True)
    args = container_args(app) + ['-d', '--workdir', app['workingDir'], '-p', '127.0.0.1::' + str(app['ports'][0]['containerPort'])]
    for env in app['env']:
        args += ['-e', env['name'] + '=' + env['value']]
    cid = run(*args, app['image'], *app['command'])
    try:
        port = int(run('podman', 'port', cid).split(':')[-1])
        for attempt in range(30):
            try:
                with urllib.request.urlopen(f'http://127.0.0.1:{port}' + app['readinessProbe']['httpGet']['path'], timeout=2) as response:
                    assert response.status == 200
                break
            except Exception:
                time.sleep(1)
        else:
            raise AssertionError('noVNC did not start')
        # A real WebSocket upgrade must reach the local X11 VNC server.
        with socket.create_connection(('127.0.0.1', port), timeout=10) as ws:
            key = base64.b64encode(os.urandom(16)).decode()
            ws.sendall((f'GET /websockify HTTP/1.1\r\nHost: localhost:{port}\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: 13\r\nSec-WebSocket-Protocol: binary\r\n\r\n').encode())
            received = b''
            while b'RFB ' not in received:
                data = ws.recv(4096)
                assert data, 'VNC handshake closed early'
                received += data
            assert b'101 Switching Protocols' in received
        listeners = run('podman', 'exec', cid, 'netstat', '-lntp')
        vnc_listeners = [line.split()[3] for line in listeners.splitlines() if 'x11vnc' in line]
        assert vnc_listeners and all(address.startswith('127.0.0.1:') for address in vnc_listeners), listeners
        status = run('podman', 'exec', cid, 'cat', '/proc/1/status')
        assert 'NoNewPrivs:\t1' in status and 'CapEff:\t0000000000000000' in status
        assert run('podman', 'exec', cid, 'id', '-u') != '0'
        processes = run('podman', 'exec', cid, 'ps', '-eo', 'args')
        assert '/usr/lib/chromium-browser/chromium-browser' in processes
        assert '/usr/bin/Xvfb' in processes
        print('Non-root read-only runtime, browser, HTTP, WebSocket/RFB, and zero capabilities passed')
    finally:
        subprocess.run(['podman', 'rm', '-f', cid], check=True, stdout=subprocess.DEVNULL)
        subprocess.run(container_args(init) + ['--rm', '--network=none', init['image'], '/bin/sh', '-ec', 'rm -rf /opt/novnc/site'], check=True)