open-swe/tests/test_atlassian_connect.py

"""Atlassian Connect qsh vectors + JWT verification (auth boundary).

qsh correctness is a silent-auth-bypass surface, so the official Atlassian test
vector and the three endpoint vectors are pinned here. The JWT tests exercise
alg-pinning, signature, exp, issuer binding, and qsh binding.
"""

from __future__ import annotations

import time
from types import SimpleNamespace

import jwt
import pytest

from agent.utils import atlassian_connect as ac

_SECRET = "connect-shared-secret"


def _request(method: str, path: str, query: str = "", *, token: str | None = None) -> object:
    headers = {"Authorization": f"JWT {token}"} if token else {}
    return SimpleNamespace(
        method=method,
        url=SimpleNamespace(path=path, query=query),
        headers=headers,
        query_params={},
    )


def _make_token(
    *,
    secret: str = _SECRET,
    iss: str = "tenant-1",
    alg: str = "HS256",
    exp_delta: int = 180,
    qsh: str | None = "auto",
    method: str = "POST",
    path: str = "/connect/webhook/comment-created",
    query: str = "",
    drop_exp: bool = False,
) -> str:
    claims: dict = {"iss": iss}
    if not drop_exp:
        claims["exp"] = int(time.time()) + exp_delta
    if qsh == "auto":
        claims["qsh"] = ac.compute_qsh(method, path, query)
    elif qsh is not None:
        claims["qsh"] = qsh
    return jwt.encode(claims, secret, algorithm=alg)


# --- qsh vectors -----------------------------------------------------------


def test_official_atlassian_qsh_vector() -> None:
    canon = ac.canonical_request(
        "GET",
        "/path/to/service",
        "zee_last=param&repeated=parameter 1&first=param&repeated=parameter 2",
    )
    assert canon == (
        "GET&/path/to/service&first=param&repeated=parameter%201,parameter%202&zee_last=param"
    )


@pytest.mark.parametrize(
    ("path", "expected"),
    [
        ("/connect/installed", "72c0a77bd4d709a202e9b2561ed003fdb400318f7a1cfabe47576d1e1d5b5dd7"),
        (
            "/connect/uninstalled",
            "ef0c0673ed4cf59a823d82cdc5c397c8643d79db724ce7d567342ea15e02acfe",
        ),
        (
            "/connect/webhook/comment-created",
            "72e058a8906e894732718ec80dbbbf073640341b9ac6ed7cd86f887f23a00b4d",
        ),
    ],
)
def test_endpoint_qsh_vectors(path: str, expected: str) -> None:
    assert ac.compute_qsh("POST", path, "") == expected


def test_qsh_drops_jwt_param_and_encodes_space_not_plus() -> None:
    # jwt param is excluded; space must be %20 (never +).
    with_jwt = ac.compute_qsh("GET", "/x", "a=b c&jwt=zzz")
    without = ac.compute_qsh("GET", "/x", "a=b c")
    assert with_jwt == without
    assert "%20" in ac.canonical_request("GET", "/x", "a=b c")
    assert "+" not in ac.canonical_request("GET", "/x", "a=b c")


# --- JWT verification ------------------------------------------------------


def test_valid_token_accepted() -> None:
    token = _make_token()
    req = _request("POST", "/connect/webhook/comment-created", token=token)
    claims = ac.verify_connect_jwt(req, shared_secret=_SECRET)
    assert claims is not None
    assert claims["iss"] == "tenant-1"


def test_missing_token_rejected() -> None:
    req = _request("POST", "/connect/webhook/comment-created")
    assert ac.verify_connect_jwt(req, shared_secret=_SECRET) is None


def test_alg_none_rejected() -> None:
    token = jwt.encode({"iss": "t", "exp": int(time.time()) + 60}, "", algorithm="none")
    req = _request("POST", "/connect/webhook/comment-created", token=token)
    assert ac.verify_connect_jwt(req, shared_secret=_SECRET) is None


def test_wrong_secret_rejected() -> None:
    token = _make_token(secret="attacker-secret")
    req = _request("POST", "/connect/webhook/comment-created", token=token)
    assert ac.verify_connect_jwt(req, shared_secret=_SECRET) is None


def test_expired_token_rejected() -> None:
    token = _make_token(exp_delta=-3600)
    req = _request("POST", "/connect/webhook/comment-created", token=token)
    assert ac.verify_connect_jwt(req, shared_secret=_SECRET) is None


def test_missing_secret_fails_closed() -> None:
    token = _make_token()
    req = _request("POST", "/connect/webhook/comment-created", token=token)
    assert ac.verify_connect_jwt(req, shared_secret=None) is None


def test_issuer_binding_mismatch_rejected() -> None:
    token = _make_token(iss="tenant-1")
    req = _request("POST", "/connect/webhook/comment-created", token=token)
    assert ac.verify_connect_jwt(req, shared_secret=_SECRET, expected_client_key="tenant-2") is None


def test_qsh_mismatch_rejected_cross_endpoint_replay() -> None:
    # Token signed with the qsh for /installed, replayed at the webhook endpoint.
    token = _make_token(path="/connect/installed")
    req = _request("POST", "/connect/webhook/comment-created", token=token)
    assert ac.verify_connect_jwt(req, shared_secret=_SECRET) is None


def test_missing_qsh_rejected_when_required() -> None:
    token = _make_token(qsh=None)
    req = _request("POST", "/connect/webhook/comment-created", token=token)
    assert ac.verify_connect_jwt(req, shared_secret=_SECRET, qsh_required=True) is None


def test_missing_qsh_allowed_on_lifecycle_when_not_required() -> None:
    token = _make_token(qsh=None)
    req = _request("POST", "/connect/installed", token=token)
    assert ac.verify_connect_jwt(req, shared_secret=_SECRET, qsh_required=False) is not None


def test_context_qsh_rejected() -> None:
    token = _make_token(qsh="context-qsh")
    req = _request("POST", "/connect/webhook/comment-created", token=token)
    assert ac.verify_connect_jwt(req, shared_secret=_SECRET, qsh_required=False) is None


def test_missing_exp_rejected() -> None:
    token = _make_token(drop_exp=True)
    req = _request("POST", "/connect/webhook/comment-created", token=token)
    assert ac.verify_connect_jwt(req, shared_secret=_SECRET) is None


# --- baseUrl host allowlist (first-install gate) ---------------------------


def test_base_url_allowlist(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(
        ac, "CONNECT_EXPECTED_BASE_URL_HOSTS", frozenset({"seahaven.atlassian.net"})
    )
    assert ac.base_url_host_allowed("https://seahaven.atlassian.net/wiki") is True
    assert ac.base_url_host_allowed("https://evil.example.com/wiki") is False


def test_base_url_allowlist_empty_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(ac, "CONNECT_EXPECTED_BASE_URL_HOSTS", frozenset())
    assert ac.base_url_host_allowed("https://seahaven.atlassian.net/wiki") is False


# --- signed-install (asymmetric RS256) lifecycle verification ---------------

_AUD = "https://openswe.example.com"


def _rsa_keypair() -> tuple[str, str]:
    from cryptography.hazmat.primitives import serialization
    from cryptography.hazmat.primitives.asymmetric import rsa

    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    priv = key.private_bytes(
        serialization.Encoding.PEM,
        serialization.PrivateFormat.PKCS8,
        serialization.NoEncryption(),
    ).decode()
    pub = (
        key.public_key()
        .public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo)
        .decode()
    )
    return priv, pub


def _install_token(priv_pem: str, *, iss="tenant-1", aud=_AUD, exp_delta=180) -> str:
    return jwt.encode(
        {"iss": iss, "aud": aud, "exp": int(time.time()) + exp_delta},
        priv_pem,
        algorithm="RS256",
        headers={"kid": "install-key-1"},
    )


def _install_req(token: str) -> object:
    return SimpleNamespace(
        method="POST",
        url=SimpleNamespace(path="/connect/installed", query=""),
        headers={"Authorization": f"JWT {token}"},
        query_params={},
    )


def _run_install_verify(token, pub_pem, monkeypatch, *, expected_client_key=None):
    import asyncio
    from unittest.mock import AsyncMock, patch

    monkeypatch.setattr(ac, "CONNECT_BASE_URL", _AUD)
    with patch.object(ac, "_fetch_atlassian_public_key", new=AsyncMock(return_value=pub_pem)):
        return asyncio.run(
            ac.verify_asymmetric_install_jwt(
                _install_req(token), expected_client_key=expected_client_key
            )
        )


def test_signed_install_valid_accepted(monkeypatch: pytest.MonkeyPatch) -> None:
    priv, pub = _rsa_keypair()
    claims = _run_install_verify(_install_token(priv), pub, monkeypatch)
    assert claims is not None and claims["iss"] == "tenant-1"


def test_signed_install_hs256_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
    # An HS256 token (symmetric alg-confusion) must not pass asymmetric verify.
    _priv, pub = _rsa_keypair()
    hs = jwt.encode({"iss": "t", "aud": _AUD, "exp": int(time.time()) + 60}, "x", algorithm="HS256")
    assert _run_install_verify(hs, pub, monkeypatch) is None


def test_signed_install_wrong_audience_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
    priv, pub = _rsa_keypair()
    token = _install_token(priv, aud="https://some-other-app.example.com")
    assert _run_install_verify(token, pub, monkeypatch) is None


def test_signed_install_wrong_key_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
    priv, _pub = _rsa_keypair()
    _priv2, pub2 = _rsa_keypair()
    assert _run_install_verify(_install_token(priv), pub2, monkeypatch) is None


def test_signed_install_expired_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
    priv, pub = _rsa_keypair()
    assert _run_install_verify(_install_token(priv, exp_delta=-3600), pub, monkeypatch) is None


def test_signed_install_issuer_binding(monkeypatch: pytest.MonkeyPatch) -> None:
    priv, pub = _rsa_keypair()
    token = _install_token(priv, iss="tenant-1")
    assert _run_install_verify(token, pub, monkeypatch, expected_client_key="tenant-2") is None


def test_signed_install_key_fetch_failure_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
    priv, _pub = _rsa_keypair()
    assert _run_install_verify(_install_token(priv), None, monkeypatch) is None


async def test_fetch_public_key_rejects_malformed_kid() -> None:
    # Defense-in-depth: a kid with path/URL chars is rejected before any fetch.
    for bad in ["../../evil", "a/b", "http://evil.com", "a b", ""]:
        assert await ac._fetch_atlassian_public_key(bad) is None


def test_client_key_allowed_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(ac, "CONNECT_EXPECTED_CLIENT_KEYS", frozenset())
    assert ac.client_key_allowed("anything") is False
    monkeypatch.setattr(ac, "CONNECT_EXPECTED_CLIENT_KEYS", frozenset({"ours"}))
    assert ac.client_key_allowed("ours") is True
    assert ac.client_key_allowed("attacker") is False
    assert ac.client_key_allowed("") is False


def test_signed_install_no_base_url_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
    import asyncio
    from unittest.mock import AsyncMock, patch

    priv, pub = _rsa_keypair()
    monkeypatch.setattr(ac, "CONNECT_BASE_URL", "")
    with patch.object(ac, "_fetch_atlassian_public_key", new=AsyncMock(return_value=pub)):
        result = asyncio.run(ac.verify_asymmetric_install_jwt(_install_req(_install_token(priv))))
    assert result is None