#!/usr/bin/python3
"""Installed root-owned pkexec helper. Never execute a user-owned script/package unchecked."""
import base64
import fcntl
import hashlib
import os
import re
import stat
import subprocess
import sys
import tempfile
from pathlib import Path

PUBLIC_KEY = "fq8ch4hdAZSsi6U3/DvftIZQWjCbEMYdcuW8RG3R1jo="
MAX_SIZE = 512 * 1024 * 1024
ENV = {"PATH": "/usr/sbin:/usr/bin:/sbin:/bin", "HOME": "/root", "LC_ALL": "C"}


def require(condition, code):
    if not condition:
        raise ValueError(code)


def validate_request(version, size, digest, signature):
    require(re.fullmatch(r"[0-9]+\.[0-9]+\.[0-9]+", version), "invalid_version")
    require(0 < size <= MAX_SIZE, "invalid_size")
    require(re.fullmatch(r"[a-f0-9]{64}", digest), "invalid_digest")
    decoded = base64.b64decode(signature, validate=True)
    require(len(decoded) == 64, "invalid_signature")
    return decoded


def copy_verified_source(source, target, uid, size, digest):
    # Open exactly once; reject links/devices and other users' files, then pin
    # the bytes in root-private staging before verification or dpkg access.
    fd = os.open(source, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK)
    with os.fdopen(fd, "rb") as inp:
        meta = os.fstat(inp.fileno())
        require(stat.S_ISREG(meta.st_mode) and meta.st_uid == uid, "invalid_source_owner_or_type")
        require(meta.st_size == size, "size_mismatch")
        count = 0
        sha = hashlib.sha256()
        with open(target, "xb") as out:
            while True:
                data = inp.read(min(65536, size - count + 1))
                if not data:
                    break
                count += len(data)
                require(count <= size, "source_grew")
                sha.update(data)
                out.write(data)
            out.flush()
            os.fsync(out.fileno())
        require(count == size and sha.hexdigest() == digest, "digest_mismatch")


def run(args, timeout=30):
    return subprocess.run(args, check=True, capture_output=True, text=True,
                          env=ENV, timeout=timeout).stdout.strip()


def verify_signature(stage, digest, signature):
    (stage / "digest").write_bytes(bytes.fromhex(digest))
    (stage / "signature").write_bytes(signature)
    (stage / "public.der").write_bytes(bytes.fromhex("302a300506032b6570032100") + base64.b64decode(PUBLIC_KEY))
    run(["/usr/bin/openssl", "pkeyutl", "-verify", "-pubin", "-keyform", "DER",
         "-inkey", str(stage / "public.der"), "-rawin", "-in", str(stage / "digest"),
         "-sigfile", str(stage / "signature")])


def install(source, version, size, digest, signature):
    require(os.geteuid() == 0, "root_required")
    uid = int(os.environ.get("PKEXEC_UID", "-1"))
    require(uid >= 0, "pkexec_required")
    sig = validate_request(version, size, digest, signature)
    os.umask(0o077)
    lock_fd = os.open("/run/soondesk-update.lock", os.O_CREAT | os.O_RDWR | os.O_NOFOLLOW, 0o600)
    with os.fdopen(lock_fd, "w") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        with tempfile.TemporaryDirectory(prefix="soondesk-update-", dir="/var/tmp") as staging:
            stage = Path(staging)
            package = stage / "update.deb"
            copy_verified_source(source, package, uid, size, digest)
            verify_signature(stage, digest, sig)
            require(run(["/usr/bin/dpkg-deb", "-f", str(package), "Package"]) == "wedesk", "wrong_package")
            require(run(["/usr/bin/dpkg-deb", "-f", str(package), "Version"]) == version, "wrong_version")
            require(run(["/usr/bin/dpkg-deb", "-f", str(package), "Architecture"]) == "amd64", "wrong_arch")
            require(run(["/usr/bin/dpkg", "--print-architecture"]) == "amd64", "unsupported_arch")
            installed = run(["/usr/bin/dpkg-query", "-W", "-f=${Version}", "wedesk"])
            run(["/usr/bin/dpkg", "--compare-versions", version, "gt", installed])
            # dry-run detects package manager locks/dependency problems before
            # allowing maintainer scripts to stop the service.
            run(["/usr/bin/dpkg", "--no-act", "--install", str(package)])
            # Keep installation output root-only for recovery diagnostics.
            log_fd = os.open("/var/log/soondesk-update.log", os.O_WRONLY | os.O_CREAT | os.O_APPEND | os.O_NOFOLLOW, 0o600)
            with os.fdopen(log_fd, "a") as log:
                log.write(f"\nSoonDesk update {installed} -> {version} sha256={digest}\n")
                log.flush()
                # Do not forcibly kill dpkg during a package transaction.
                result = subprocess.run(["/usr/bin/dpkg", "--install", str(package)],
                                        env=ENV, stdout=log, stderr=log)
                require(result.returncode == 0, "installation_failed_see_root_log")
            require(run(["/usr/bin/dpkg-query", "-W", "-f=${Version}", "wedesk"]) == version, "installed_version_mismatch")
            run(["/usr/bin/systemctl", "is-active", "wedesk"])
    print("installed")


if __name__ == "__main__":
    try:
        require(len(sys.argv) == 6, "invalid_arguments")
        install(sys.argv[1], sys.argv[2], int(sys.argv[3]), sys.argv[4], sys.argv[5])
    except Exception as error:
        # Avoid printing package contents, credentials or arbitrary command output.
        print(f"SoonDesk update failed: {type(error).__name__}: " +
              (str(error) if isinstance(error, ValueError) else "authorization_or_validation_or_installation_failed"), file=sys.stderr)
        sys.exit(1)
