Coverage for pki_tools/types/utils.py: 84%

1import binascii 

2import socket 

3import ssl 

4from functools import cache 

5from urllib.parse import urlparse 

6 

7import httpx 

8from loguru import logger 

9from pydantic import BaseModel, constr 

10 

11from pki_tools.exceptions import FetchFailure 

12 

13CACHE_TIME_SECONDS = 60 * 60 * 24 * 30 # 1 month 

14HTTPX_CLIENT = httpx.Client( 

15 transport=httpx.HTTPTransport(retries=2), timeout=30 

16) 

17 

18 

19class CertsUri(BaseModel): 

20 """ 

21 Describes a URI where one or more public certificate(s) 

22 can be downloaded 

23 

24 Attributes: 

25 uri: The URI for the public certificate(s) 

26 cache_time_seconds: Specifies how long the public cert should be 

27 cached, default is 1 month. 

28 """ 

29 

30 uri: constr(pattern=r"https*://.*") 

31 cache_time_seconds: int = CACHE_TIME_SECONDS 

32 

33 @property 

34 def hostname(self): 

35 return urlparse(self.uri).hostname 

36 

37 

38def _byte_to_hex(bytes_in: bytes) -> str: 

39 return binascii.hexlify(bytes_in).decode().upper() 

40 

41 

42def _hex_to_byte(hex_string: str) -> bytes: 

43 byte_array = bytearray() 

44 for i in range(0, len(hex_string), 2): 

45 byte_array.append(int(hex_string[i : i + 2], 16)) 

46 

47 return bytes(byte_array) 

48 

49 

50@cache 

51def _download_server_certificate(hostname: str, cache_ttl: int | None = None): 

52 context = ssl.create_default_context() 

53 with ( 

54 socket.create_connection((hostname, 443)) as sock, 

55 context.wrap_socket(sock, server_hostname=hostname) as ssock, 

56 ): 

57 der_cert = ssock.getpeercert(binary_form=True) 

58 return ssl.DER_cert_to_PEM_cert(der_cert) 

59 

60 

61def _do_download(client: httpx.Client, uri: str) -> httpx.Response: 

62 """Helper to perform the actual download and check status.""" 

63 ret = client.get(uri) 

64 

65 if ret.status_code != 200: 

66 logger.bind(status=ret.status_code).error( 

67 "Failed to fetch issuer from URI" 

68 ) 

69 raise FetchFailure(f"Failed to fetch URI. Status: {ret.status_code}") 

70 

71 return ret 

72 

73 

74@cache 

75def _download_cached( 

76 uri: str, ttl: int | None = None, proxy: str | None = None 

77) -> httpx.Response: 

78 """ 

79 Download and cache content from a URI with optional proxy support. 

80 

81 Args: 

82 uri: The URI to download from 

83 ttl: Time-to-live for cache (used as cache key for cache invalidation) 

84 proxy: Optional proxy URL (e.g., 'http://proxy.example.com:8080') 

85 

86 Returns: 

87 httpx.Response object containing the downloaded content 

88 """ 

89 # Use cached version when proxy is None 

90 if proxy is None: 

91 client = HTTPX_CLIENT 

92 else: 

93 # For proxy requests, create a new client with proxy 

94 client = httpx.Client( 

95 transport=httpx.HTTPTransport(retries=2), timeout=30, proxy=proxy 

96 ) 

97 

98 try: 

99 return _do_download(client, uri) 

100 finally: 

101 if proxy is not None: 

102 client.close()