#!/usr/bin/env python3
"""Own-creator OTA signed client. Uses only the Python standard library.

Real device keys are read from a UTF-8 file or environment variable, never
printed. The shipped vector key is public test data and is not a device key.
"""
from __future__ import annotations

import argparse
import base64
import hashlib
import hmac
import json
import os
import secrets
import sys
import tempfile
import time
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.parse import urlsplit
from urllib.request import HTTPRedirectHandler, Request, build_opener

MAX_APPLICATION_BYTES = 8 * 1024 * 1024


def canonical_request(method, raw_path, raw_body, serial_number, timestamp, nonce, activation_id=""):
    values = (method, raw_path, serial_number, str(timestamp), nonce, activation_id)
    if any("\n" in str(value) or "\r" in str(value) for value in values):
        raise ValueError("Signed request fields cannot contain newlines")
    return "\n".join(("request:v1", f"method:{method.upper()}", f"path:{raw_path}",
        f"body_sha256:{hashlib.sha256(raw_body).hexdigest()}", f"serial_number:{serial_number}",
        f"timestamp:{timestamp}", f"nonce:{nonce}", f"activation_id:{activation_id}"))


def request_signature(key, method, raw_path, raw_body, serial_number, timestamp, nonce, activation_id=""):
    canonical = canonical_request(method, raw_path, raw_body, serial_number, timestamp, nonce, activation_id)
    digest = hmac.new(key.encode("utf-8"), canonical.encode("utf-8"), hashlib.sha256).digest()
    return base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=")


def signed_headers(*, key, method, raw_path, raw_body, serial_number, device_id,
                   hardware_version, firmware_version, activation_id="", client_id="creator-ota-cli",
                   timestamp=None, nonce=None):
    timestamp = str(int(time.time()) if timestamp is None else timestamp)
    nonce = nonce or secrets.token_hex(16)
    return {"Serial-Number": serial_number, "Device-Id": device_id,
        "Hardware-Version": hardware_version, "Firmware-Version": firmware_version,
        "Activation-Id": activation_id, "Client-Id": client_id,
        "X-Device-Timestamp": timestamp, "X-Device-Nonce": nonce,
        "X-Device-Signature-Alg": "HMAC-SHA256",
        "X-Device-Signature": request_signature(key, method, raw_path, raw_body,
            serial_number, timestamp, nonce, activation_id),
        "Product": "own_creator", "Accept-Encoding": "identity"}


class NoRedirects(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def safe_response(value):
    """Printable response keeps credentials out of shell logs."""
    if isinstance(value, dict):
        return {key: "<redacted>" if key.lower() in {"token", "password", "device_key", "signature", "authorization", "access_token", "x-device-signature"}
                else safe_response(item) for key, item in value.items()}
    if isinstance(value, list):
        return [safe_response(item) for item in value]
    if isinstance(value, str) and value.startswith(("http://", "https://", "ws://", "wss://")):
        return value.split("?", 1)[0] + ("?<redacted>" if "?" in value else "")
    return value


def verify_vectors(path):
    document = json.loads(Path(path).read_text(encoding="utf-8"))
    for vector in document["vectors"]:
        raw = vector["body_utf8"].encode("utf-8")
        args = (vector["method"], vector["raw_path"], raw, vector["serial_number"],
                vector["timestamp"], vector["nonce"], vector["activation_id"])
        if canonical_request(*args) != vector["canonical"]:
            raise ValueError(f"Canonical vector failed: {vector['name']}")
        if request_signature(document["public_test_key"], *args) != vector["signature"]:
            raise ValueError(f"Signature vector failed: {vector['name']}")
    return len(document["vectors"])


def default_vectors_path():
    script = Path(__file__).resolve()
    candidates = [script.with_name("signature-vectors.json"), script.with_name("creator-ota-signature-vectors.json")]
    if len(script.parents) > 3:
        candidates.append(script.parents[3] / "docs/creator-ota-signature-vectors.json")
    for path in candidates:
        if path.is_file():
            return path
    raise ValueError("Download signature-vectors.json and supply verify-vectors --file /path/to/signature-vectors.json")


def read_key(args):
    if args.key_file:
        key = Path(args.key_file).read_text(encoding="utf-8").rstrip("\r\n")
    else:
        key = os.environ.get(args.key_env or "CREATOR_OTA_DEVICE_KEY", "")
    if not key:
        raise ValueError("Provide a device key with --key-file or the selected environment variable")
    return key


def download(response, output, expected_size, expected_sha):
    if response.status != 200:
        raise ValueError("Firmware download must return HTTP 200")
    if response.headers.get("Content-Encoding", "identity").lower() not in {"", "identity"}:
        raise ValueError("Compressed firmware response rejected")
    if response.headers.get("Transfer-Encoding"):
        raise ValueError("Chunked firmware response rejected")
    if response.headers.get("Content-Length") != str(expected_size):
        raise ValueError("Firmware Content-Length does not match the checked manifest")
    destination = Path(output)
    digest, received = hashlib.sha256(), 0
    temporary = None
    try:
        with tempfile.NamedTemporaryFile(dir=destination.parent, prefix=".creator-download-", delete=False) as stream:
            temporary = Path(stream.name)
            while True:
                chunk = response.read(min(65536, expected_size + 1 - received))
                if not chunk:
                    break
                received += len(chunk)
                if received > expected_size:
                    raise ValueError("Firmware exceeds the checked size")
                digest.update(chunk)
                stream.write(chunk)
            if received != expected_size or not hmac.compare_digest(digest.hexdigest(), expected_sha.lower()):
                raise ValueError("Firmware length or SHA-256 does not match the checked manifest")
            stream.flush()
            os.fsync(stream.fileno())
        # Atomic create without replacing an existing user file.
        os.link(temporary, destination)
    finally:
        if temporary is not None:
            temporary.unlink(missing_ok=True)
    return {"downloaded": True, "bytes": received, "sha256": digest.hexdigest(),
            "installation_confirmed": False}


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    verify = commands.add_parser("verify-vectors", help="Check the public offline signature fixtures")
    verify.add_argument("--file", help="Public vector JSON; also finds a sibling downloaded signature-vectors.json")
    for name in ("check", "bootstrap", "download"):
        sub = commands.add_parser(name)
        sub.add_argument("--url", required=True, help="Exact URL, including any signed query string")
        sub.add_argument("--serial-number", required=True)
        sub.add_argument("--device-id", required=True)
        sub.add_argument("--hardware-version", required=True)
        sub.add_argument("--firmware-version", required=True)
        sub.add_argument("--activation-id", default="")
        sub.add_argument("--client-id", default="creator-ota-cli")
        keys = sub.add_mutually_exclusive_group()
        keys.add_argument("--key-file", help="UTF-8 file containing the actual device key")
        keys.add_argument("--key-env", help="Environment variable name; defaults to CREATOR_OTA_DEVICE_KEY")
        sub.add_argument("--allow-http", action="store_true", help="Explicitly allow controlled local HTTP testing")
        if name == "download":
            sub.add_argument("--output", required=True, help="New output file; existing files are never replaced")
            sub.add_argument("--sha256", required=True)
            sub.add_argument("--size", required=True, type=int)
        else:
            sub.add_argument("--body-file", required=True, help="Signed verbatim, without JSON reserialization")
    args = parser.parse_args(argv)
    if args.command == "verify-vectors":
        print(json.dumps({"vectors_verified": verify_vectors(args.file or default_vectors_path())}))
        return 0
    parsed = urlsplit(args.url)
    if parsed.scheme not in {"http", "https"} or not parsed.hostname or parsed.username or parsed.password or parsed.fragment:
        raise ValueError("Use a complete HTTP(S) URL without user credentials or fragments")
    if parsed.scheme == "http" and not args.allow_http:
        raise ValueError("Use HTTPS, or explicitly pass --allow-http for isolated local testing")
    raw_path = (parsed.path or "/") + ("?" + parsed.query if parsed.query else "")
    raw = b"" if args.command == "download" else Path(args.body_file).read_bytes()
    if args.command != "download":
        body = json.loads(raw)
        intent = "check_ota" if args.command == "check" else "bootstrap"
        if not isinstance(body, dict) or body.get("intent") != intent:
            raise ValueError(f"The raw body file must contain intent={intent}")
    elif not 0 < args.size <= MAX_APPLICATION_BYTES or len(args.sha256) != 64 or any(c not in "0123456789abcdefABCDEF" for c in args.sha256):
        raise ValueError("Download requires the candidate's positive size (up to 8 MiB) and SHA-256")
    method = "GET" if args.command == "download" else "POST"
    headers = signed_headers(key=read_key(args), method=method, raw_path=raw_path, raw_body=raw,
        serial_number=args.serial_number, device_id=args.device_id, hardware_version=args.hardware_version,
        firmware_version=args.firmware_version, activation_id=args.activation_id, client_id=args.client_id)
    if method == "POST":
        headers["Content-Type"] = "application/json"
    request = Request(args.url, data=raw if method == "POST" else None, headers=headers, method=method)
    with build_opener(NoRedirects).open(request, timeout=30) as response:
        if args.command == "download":
            result = download(response, args.output, args.size, args.sha256)
        else:
            result = safe_response(json.loads(response.read(1024 * 1024)))
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except HTTPError as error:
        # Never include error.url (which could carry a signed URL query).
        try:
            detail = safe_response(json.loads(error.read(65536)))
        except (UnicodeError, ValueError):
            detail = {"message": "HTTP request failed"}
        print(json.dumps({"status": error.code, "error": detail}, ensure_ascii=False), file=sys.stderr)
        sys.exit(1)
    except (ValueError, OSError, URLError) as error:
        # urllib errors can contain the URL. Keep network failures generic.
        text = "Network request failed" if isinstance(error, URLError) else str(error)
        print(text, file=sys.stderr)
        sys.exit(1)