Coverage for paperless/signed_pickle.py: 52%

1from __future__ import annotations 

2 

3import hmac 

4import pickle 

5from hashlib import sha256 

6from typing import Any 

7 

8from django.conf import settings 

9 

10HMAC_SIZE = sha256().digest_size 

11 

12 

13class SignedPickleError(ValueError): 

14 """Raised when a signed pickle payload cannot be authenticated.""" 

15 

16 

17def _get_signing_key() -> bytes: 

18 return settings.SECRET_KEY.encode() 

19 

20 

21def signed_pickle_dumps(obj: object) -> bytes: 

22 data = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) 

23 signature = hmac.new(_get_signing_key(), data, sha256).digest() 

24 return signature + data 

25 

26 

27def signed_pickle_loads(payload: bytes) -> Any: 

28 if len(payload) <= HMAC_SIZE: 

29 msg = "Signed-pickle payload too short" 

30 raise SignedPickleError(msg) 

31 

32 signature = payload[:HMAC_SIZE] 

33 data = payload[HMAC_SIZE:] 

34 expected = hmac.new(_get_signing_key(), data, sha256).digest() 

35 if not hmac.compare_digest(signature, expected): 

36 msg = "Signed-pickle HMAC verification failed; payload may have been tampered with" 

37 raise SignedPickleError(msg) 

38 

39 return pickle.loads(data)