Coverage for src/twofas/_security.py: 100%

75 statements  

« prev     ^ index     » next       coverage.py v7.4.0, created at 2024-01-22 20:51 +0100

1import base64 

2import getpass 

3import hashlib 

4import json 

5import logging 

6import time 

7import warnings 

8from pathlib import Path 

9from typing import Any, Optional 

10 

11import cryptography.exceptions 

12import keyring 

13import keyring.backends.SecretService 

14from cryptography.hazmat.primitives.ciphers.aead import AESGCM 

15from cryptography.hazmat.primitives.hashes import SHA256 

16from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC 

17from keyring.backend import KeyringBackend 

18 

19from ._types import AnyDict, TwoFactorAuthDetails, into_class 

20 

21# Suppress keyring warnings 

22keyring_logger = logging.getLogger("keyring") 

23keyring_logger.setLevel(logging.ERROR) # Set the logging level to ERROR for keyring logger 

24 

25 

26def _decrypt(encrypted: str, passphrase: str) -> list[AnyDict]: 

27 # thanks https://github.com/wodny/decrypt-2fas-backup/blob/master/decrypt-2fas-backup.py 

28 credentials_enc, pbkdf2_salt, nonce = map(base64.b64decode, encrypted.split(":")) 

29 kdf = PBKDF2HMAC(algorithm=SHA256(), length=32, salt=pbkdf2_salt, iterations=10000) 

30 key = kdf.derive(passphrase.encode()) 

31 aesgcm = AESGCM(key) 

32 credentials_dec = aesgcm.decrypt(nonce, credentials_enc, None) 

33 dec = json.loads(credentials_dec) # type: list[AnyDict] 

34 if not isinstance(dec, list): # pragma: no cover 

35 raise TypeError("Unexpected data structure in input file.") 

36 return dec 

37 

38 

39def decrypt(encrypted: str, passphrase: str) -> list[TwoFactorAuthDetails]: 

40 """ 

41 

42 Raises: 

43 PermissionError 

44 """ 

45 try: 

46 dicts = _decrypt(encrypted, passphrase) 

47 return into_class(dicts, TwoFactorAuthDetails) 

48 except cryptography.exceptions.InvalidTag as e: 

49 # wrong passphrase! 

50 raise PermissionError("Invalid passphrase for file.") from e 

51 

52 

53def hash_string(data: Any) -> str: 

54 """ 

55 Hashes a string using SHA-256. 

56 """ 

57 sha256 = hashlib.sha256() 

58 sha256.update(str(data).encode()) 

59 return sha256.hexdigest() 

60 

61 

62PREFIX = "2fas:" 

63 

64 

65class KeyringManager: 

66 appname: str = "" 

67 tmp_file = Path("/tmp") / ".2fas" 

68 

69 def __init__(self) -> None: 

70 self._init() 

71 

72 def _init(self) -> None: 

73 # so you can call init again to set active appname (for pytest) 

74 tmp_file = self.tmp_file 

75 # APPNAME is session specific but with global prefix for easy clean up 

76 

77 if tmp_file.exists() and (session := tmp_file.read_text()) and session.startswith(PREFIX): 

78 # existing session 

79 self.appname = session 

80 else: 

81 # new session! 

82 session = hash_string((time.time())) # random enough for this purpose 

83 self.appname = f"{PREFIX}{session}" 

84 tmp_file.write_text(self.appname) 

85 

86 @classmethod 

87 def _retrieve_credentials(cls, filename: str, appname: str) -> Optional[str]: 

88 return keyring.get_password(appname, hash_string(filename)) 

89 

90 def retrieve_credentials(self, filename: str) -> Optional[str]: 

91 return self._retrieve_credentials(filename, self.appname) 

92 

93 @classmethod 

94 def _save_credentials(cls, filename: str, passphrase: str, appname: str) -> None: 

95 keyring.set_password(appname, hash_string(filename), passphrase) 

96 

97 def save_credentials(self, filename: str) -> str: 

98 passphrase = getpass.getpass(f"Passphrase for '{filename}'? ") 

99 self._save_credentials(filename, passphrase, self.appname) 

100 

101 return passphrase 

102 

103 @classmethod 

104 def _delete_credentials(cls, filename: str, appname: str) -> None: 

105 keyring.delete_password(appname, hash_string(filename)) 

106 

107 def delete_credentials(self, filename: str) -> None: 

108 self._delete_credentials(filename, self.appname) 

109 

110 @classmethod 

111 def _cleanup_keyring(cls, appname: str) -> int: 

112 kr: keyring.backends.SecretService.Keyring | KeyringBackend = keyring.get_keyring() 

113 if not hasattr(kr, "get_preferred_collection"): # pragma: no cover 

114 warnings.warn(f"Can't clean up this keyring backend! {type(kr)}", category=RuntimeWarning) 

115 return -1 

116 

117 collection = kr.get_preferred_collection() 

118 

119 # get old 2fas: keyring items: 

120 return len( 

121 [ 

122 item 

123 for item in collection.get_all_items() 

124 if ( 

125 service := item.get_attributes().get("service", "") 

126 ) # must have a 'service' attribute, otherwise it's unrelated 

127 and service.startswith(PREFIX) # must be a 2fas: service, otherwise it's unrelated 

128 and service != appname # must not be the currently active session 

129 ] 

130 ) 

131 

132 def cleanup_keyring(self) -> None: 

133 self._cleanup_keyring(self.appname) 

134 

135 

136keyring_manager = KeyringManager()