from __future__ import annotations

import hashlib
import base64
import json
import unittest
import os
import re
import subprocess
import tempfile

import yaml
from pathlib import Path

from scripts.verified_l2_signing_policy import BUILDER, WORKFLOW, select

ROOT = Path(__file__).resolve().parents[1]


class VerifiedSigningPolicyTests(unittest.TestCase):
    def setUp(self):
        self.policy = json.loads((ROOT / "release/serviceification-image-signing-policy.json").read_text())
        self.legacy = {**self.policy, "schema_version": "nexuscrm.l2-oci-owner-signing-policy/v2"}
        self.legacy["trusted_workflow_ref"] = BUILDER + "b" * 40
        self.legacy_bytes = (json.dumps(self.legacy, indent=2) + "\n").encode()

    def signature(self, *, embedded=True, authority="a" * 40):
        raw = ((json.dumps(self.policy, sort_keys=True, separators=(",", ":")) + "\n").encode()
               if embedded else self.legacy_bytes)
        annotations = {
            "io.nexuscrm.source.ref": "refs/heads/original-producer",
            "io.nexuscrm.signer.workflow": WORKFLOW + ("@refs/heads/main" if embedded else "@refs/heads/original-producer"),
            "io.nexuscrm.signer.policy-sha256": "sha256:" + hashlib.sha256(raw).hexdigest(),
        }
        if embedded:
            annotations["io.nexuscrm.signer.policy-base64"] = base64.b64encode(raw).decode().rstrip("=")
            annotations["io.nexuscrm.signer.authority-commit"] = authority
        return {"optional": annotations}

    def test_workflow_annotation_encoding_is_cosign_safe(self):
        workflow = yaml.safe_load((ROOT / ".github/workflows/trusted-publish-serviceification-test-stack.yml").read_text())
        step = next(step for step in workflow["jobs"]["sign"]["steps"]
                    if step["name"] == "Sign exact B1 subject and canonical provenance")
        expression = re.search(r'-a "(io\.nexuscrm\.signer\.policy-base64=.*)" \\', step["run"]).group(1)
        with tempfile.TemporaryDirectory() as directory:
            policy = Path(directory) / "policy.json"
            for raw in (b"{}", b"{}\n", b"{ }\n"):
                policy.write_bytes(raw)
                result = subprocess.run(["bash", "-c", 'printf "%s" "' + expression + '"'],
                                        env={**os.environ, "SIGNING_POLICY": str(policy)},
                                        text=True, capture_output=True, check=True)
                annotation = result.stdout
                self.assertEqual(1, annotation.count("="))
                self.assertNotRegex(annotation, r'[\s,"\']')
                encoded = annotation.split("=", 1)[1]
                self.assertEqual(raw, base64.b64decode(encoded + "=" * (-len(encoded) % 4), validate=True))

    def test_signature_policy_is_independent_of_producer_policy_and_branch(self):
        raw, selected, context = select([self.signature()], b'{}')
        self.assertEqual(self.policy, json.loads(raw))
        self.assertEqual("refs/heads/original-producer", context["source_ref"])
        self.assertEqual(WORKFLOW + "@refs/heads/main", context["workflow_ref"])
        self.assertEqual("a" * 40, context["authority_commit"])

    def test_legacy_release_still_requires_its_original_policy(self):
        signature = self.signature(embedded=False)
        raw, _, context = select([signature], self.legacy_bytes)
        self.assertEqual(self.legacy_bytes, raw)
        self.assertEqual("", context["authority_commit"])
        with self.assertRaisesRegex(ValueError, "no verified L2 signature"):
            select([signature], b'{}')

    def test_policy_only_resigning_accepts_multiple_valid_signatures(self):
        signatures = [self.signature(authority="b" * 40), self.signature(), self.signature(embedded=False)]
        selected = select(signatures, self.legacy_bytes)
        self.assertEqual(selected, select(list(reversed(signatures)), self.legacy_bytes))
        self.assertEqual("a" * 40, selected[2]["authority_commit"])

    def test_tampered_policy_is_not_accepted(self):
        signature = self.signature()
        signature["optional"]["io.nexuscrm.signer.policy-base64"] = base64.b64encode(b"{}").decode()
        with self.assertRaisesRegex(ValueError, "no verified L2 signature"):
            select([signature], self.legacy_bytes)

    def test_unreviewed_branch_and_wrong_owner_are_rejected(self):
        for field, value in (("trusted_workflow_ref", BUILDER + "refs/heads/feature"),
                             ("owner_repository", "unrelated/repository")):
            with self.subTest(field=field):
                self.policy[field] = value
                with self.assertRaises(ValueError):
                    select([self.signature()], self.legacy_bytes)
                self.setUp()

    def test_historical_pinned_signer_remains_verifiable(self):
        self.policy["trusted_workflow_ref"] = BUILDER + "a" * 40
        _, _, context = select([self.signature()], self.legacy_bytes)
        self.assertEqual(BUILDER + "a" * 40, context["builder_ref"])

    def test_missing_authority_commit_is_rejected(self):
        signature = self.signature()
        del signature["optional"]["io.nexuscrm.signer.authority-commit"]
        with self.assertRaisesRegex(ValueError, "authority revision"):
            select([signature], self.legacy_bytes)

    def test_duplicate_policy_fields_are_rejected(self):
        signature = self.signature()
        signature["optional"]["io.nexuscrm.signer.policy-base64"] = base64.b64encode(b'{"x":1,"x":2}').decode()
        with self.assertRaisesRegex(ValueError, "duplicate signing-policy field"):
            select([signature], self.legacy_bytes)


if __name__ == "__main__":
    unittest.main()