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

1import argparse 3 ctx1abc

2import pathlib 3 ctx1abc

3import sys 3 ctx1abc

4from typing import Annotated, Any 3 ctx1abc

5 

6from authlib.jose import JsonWebKey, KeySet 3 ctx1abc

7from pydantic import PlainValidator 3 ctx1abc

8 

9 

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

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

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

13 keyset = KeySet(keys=[key]) 1 ctx1a

14 return keyset.as_json(is_private=True) 1 ctx1a

15 

16 

17TIP_MESSAGE = ( 3 ctx1abc

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: 3 ctx1abc

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

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 true2 ctx1bc

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: 2 ctx1bc

33 with open(value) as f: 2 ctx1bc

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

35 return JsonWebKey.import_key_set(content) 2 ctx1bc

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)] 3 ctx1abc

43 

44if __name__ == "__main__": 3 ctx1abc

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

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

47 parser.add_argument( 1 ctx1a

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

49 ) 

50 args = parser.parse_args() 1 ctx1a

51 

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

53 sys.stdout.write(jwks) 1 ctx1a

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