#!/usr/bin/python3

import array
import json
import shutil
import os
import socket
import subprocess
import sys


# copied from remoteloop.py
def load_fds(sock, msglen):
    fds = array.array("i")   # Array of ints
    msg, ancdata, _, addr = sock.recvmsg(msglen, socket.CMSG_LEN(253 * fds.itemsize))
    for cmsg_level, cmsg_type, cmsg_data in ancdata:
        if (cmsg_level == socket.SOL_SOCKET and cmsg_type == socket.SCM_RIGHTS):
            # Append data, ignoring any truncated integers at the end.
            fds.frombytes(cmsg_data[:len(cmsg_data) - (len(cmsg_data) % fds.itemsize)])
    return json.loads(msg), list(fds), addr

def ldconfig():
    # ld.so.conf must exist, or `ldconfig` throws a warning
    subprocess.run(["touch", "/etc/ld.so.conf"], check=True)
    subprocess.run(["ldconfig"], check=True)


def sysusers():
    try:
        subprocess.run(["systemd-sysusers"], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, check=True)
    except subprocess.CalledProcessError as error:
        sys.stderr.write(error.stdout)
        sys.exit(1)


def update_ca_trust():
    if not shutil.which("update-ca-trust"):
        return

    # generate /etc/pki/tls/certs/ca-bundle.crt
    os.makedirs("/etc/pki/ca-trust/extracted/pem")
    os.makedirs("/etc/pki/tls/certs")
    os.symlink("/etc/pki/ca-trust/extracted/pem/tls-ca-bundle.pem", "/etc/pki/tls/certs/ca-bundle.crt")

    # allow to fail, because it sometimes mysteriously does
    subprocess.run(["update-ca-trust", "extract"], check=False)


def append_certs(cert_conf, dir_fd, parents=b""):
    for entry in os.scandir(f"/proc/self/fd/{dir_fd}".encode()):
        if entry.is_file():
            line = os.path.join(parents, entry.name)
            cert_conf.write(line)
            cert_conf.write(b"\n")
        elif entry.is_dir():
            append_certs(cert_conf,
                         os.open(entry.name, os.O_DIRECTORY, dir_fd=dir_fd),
                         os.path.join(parents, entry.name))


def update_ca_certificates():
    if not shutil.which("update-ca-certificates"):
        return

    # generate /etc/ssl/certs/ca-certificates.crt
    os.makedirs("/etc/ssl/certs")
    with open("/etc/ca-certificates.conf", "wb") as f:
        append_certs(f, os.open("/usr/share/ca-certificates", os.O_DIRECTORY))
    subprocess.run(["update-ca-certificates"], check=True)


def tmpfiles():
    # Allow systemd-tmpfiles to return non-0. Some packages want to create
    # directories owned by users that are not set up with systemd-sysusers.
    subprocess.run(["systemd-tmpfiles", "--create"], check=False)


def nsswitch():
    # the default behavior is fine, but using nss-resolve does not
    # necessarily work in a non-booted container, so make sure that
    # is not configured.
    try:
        os.remove("/etc/nsswitch.conf")
    except FileNotFoundError:
        pass


def setup_stdio():
    with socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) as sock:
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_PASSCRED, 1)
        sock.connect("/run/osbuild/api/osbuild")
        req = {'method': 'setup-stdio'}
        sock.send(json.dumps(req).encode('utf-8'))
        msg, fds, _ = load_fds(sock, 1024)
        for io in ['stdin', 'stdout', 'stderr']:
            target = getattr(sys, io)
            source = fds[msg[io]]
            os.dup2(source, target.fileno())
            os.close(source)


if __name__ == "__main__":
    setup_stdio()
    ldconfig()
    sysusers()
    update_ca_trust()
    update_ca_certificates()
    tmpfiles()
    nsswitch()

    r = subprocess.run(sys.argv[1:], check=False)
    sys.exit(r.returncode)