Coverage for polar/kit/jwk.py: 89%

1import argparse 

2import pathlib 

3import sys 

4from typing import Annotated, Any 

5 

6from authlib.jose import JsonWebKey, KeySet 

7from pydantic import PlainValidator 

8 

9 

10def generate_jwks(kid: str, size: int = 2048) -> str: 

11 options = {"kid": kid, "use": "sig"} 

12 key = JsonWebKey.generate_key("RSA", size, options, is_private=True) 

13 keyset = KeySet(keys=[key]) 

14 return keyset.as_json(is_private=True) 

15 

16 

17TIP_MESSAGE = ( 

18 "If you're in local development, you can generate a JWKS file " 

19 "by running the following command:\n" 

20 "uv run task generate_dev_jwks" 

21) 

22 

23 

24def _validate_jwks(value: Any) -> KeySet: 

25 path = pathlib.Path(str(value)) 

26 if not path.exists() and not path.is_file(): 26 ↛ 27line 26 didn't jump to line 27 because the condition on line 26 was never true

27 raise ValueError( 

28 f"The provided JWKS path {value} is not a valid file path " 

29 f"or does not exist.\n{TIP_MESSAGE}" 

30 ) 

31 

32 try: 

33 with open(value) as f: 

34 content = f.read().strip() 

35 return JsonWebKey.import_key_set(content) 

36 except ValueError as e: 

37 raise ValueError( 

38 f"The provided JWKS file {value} is not a valid JWKS file.\n{TIP_MESSAGE}" 

39 ) from e 

40 

41 

42JWKSFile = Annotated[KeySet, PlainValidator(_validate_jwks)] 

43 

44if __name__ == "__main__": 

45 parser = argparse.ArgumentParser(description="Generate JWKS") 

46 parser.add_argument("kid", type=str, help="Key ID") 

47 parser.add_argument( 

48 "--size", type=int, default=2048, help="Key size (default: 2048)" 

49 ) 

50 args = parser.parse_args() 

51 

52 jwks = generate_jwks(args.kid, args.size) 

53 sys.stdout.write(jwks) 

54 sys.stdout.write("\n")