Coverage for kryptic_cypher/cypher/base.py: 100%

1""" 

2Module that contains the code for cyphers in our system. 

3 

4Defines both the interfaces as well as a registery decorator to register classes as cyphers. 

5 

6```python 

7from kryptic_cypher.cypher import register_cypher, Cypher, CypherWithKey 

8 

9@register_cypher 

10class MyCypher(Cypher): 

11 def encode(self, text: str) -> str: 

12 # Do my encryption here 

13 return text 

14 

15 def decode(self, text: str) -> str: 

16 # Do my decryption here 

17 return text 

18 

19 

20@register_cypher 

21class MyCypherWithKey(CypherWithKey): 

22 def encode(self, text: str, key: str) -> str: 

23 # Do my encryption here 

24 return text 

25 

26 def decode(self, text: str, key: str) -> str: 

27 # Do my decryption here 

28 return text 

29``` 

30 

31You can use this if you want to import and leverage any existing cyphers that have been registered with us. 

32""" 

33 

34from abc import ABC, abstractmethod 

35from logging import getLogger 

36from pydantic import BaseModel 

37 

38logger = getLogger(__name__) 

39 

40 

41class ValidationResult(BaseModel): 

42 success: bool 

43 messages: list[str] 

44 

45 @classmethod 

46 def ok(cls): 

47 return cls( 

48 success=True, 

49 messages=[], 

50 ) 

51 

52 @classmethod 

53 def fail( 

54 cls, 

55 *messages: list[str], 

56 ): 

57 return cls( 

58 success=False, 

59 messages=list(messages), 

60 ) 

61 

62 

63class CypherResult(BaseModel): 

64 original_text: str | bytes 

65 new_text: str | bytes 

66 success: bool 

67 error: str | None 

68 

69 @classmethod 

70 def ok( 

71 cls, 

72 original_text: str | bytes, 

73 new_text: str | bytes, 

74 ): 

75 return cls( 

76 original_text=original_text, 

77 new_text=new_text, 

78 success=True, 

79 error=None, 

80 ) 

81 

82 @classmethod 

83 def fail( 

84 cls, 

85 original_text: str | bytes, 

86 error: str, 

87 ): 

88 return cls( 

89 original_text=original_text, 

90 new_text="" if isinstance(original_text, str) else b"", 

91 success=False, 

92 error=error, 

93 ) 

94 

95 

96class Cypher(ABC): 

97 """A Cypher that does not require a key to encode/decode.""" 

98 

99 @classmethod 

100 def get_name(cls) -> str: # pragma: no cover 

101 """ 

102 Get the name of the Cypher. 

103 

104 **Returns** 

105 - str: The name of the Cypher 

106 """ 

107 return cls.__module__.split(".")[-1] 

108 

109 @abstractmethod 

110 def encode(self, text: str | bytes) -> CypherResult: # pragma: no cover 

111 """Encode the given text using the given key. 

112 

113 **Parameters** 

114 - text (str): The text to encode 

115 

116 **Returns** 

117 - IO: The encoded text 

118 """ 

119 pass 

120 

121 @abstractmethod 

122 def decode(self, text: str | bytes) -> CypherResult: # pragma: no cover 

123 """Decode the given text using the given key. 

124 

125 **Parameters** 

126 - text (str): The text to decode 

127 

128 **Returns** 

129 - IO: The decoded text 

130 """ 

131 pass 

132 

133 

134class CypherWithKey(ABC): 

135 """A Cypher that requires a key to encode/decode.""" 

136 

137 @classmethod 

138 def get_name(cls) -> str: # pragma: no cover 

139 """ 

140 Get the name of the Cypher. 

141 

142 **Returns** 

143 - str: The name of the Cypher 

144 """ 

145 return cls.__module__.split(".")[-1] 

146 

147 @classmethod 

148 @abstractmethod 

149 def validate_key(cls, key: str) -> ValidationResult: # pragma: no cover 

150 """ 

151 Validate the key for the Cypher. If the key is invalid, it should return a ValidationResult 

152 with a message explaining why the key is invalid. 

153 

154 **Parameters** 

155 - key (str): The key to validate 

156 

157 **Returns** 

158 - ValidationResult: ValidationResult indicating if the key is valid or not 

159 """ 

160 pass 

161 

162 @abstractmethod 

163 def encode( 

164 self, text: str | bytes, key: str | bytes 

165 ) -> CypherResult: # pragma: no cover 

166 """Encode the given text using the given key. 

167 

168 **Parameters** 

169 - text (str): The text to encode 

170 - key (str): The key to use 

171 

172 **Returns** 

173 - IO: The encoded text 

174 """ 

175 pass 

176 

177 @abstractmethod 

178 def decode( 

179 self, 

180 text: str | bytes, 

181 key: str | bytes, 

182 ) -> CypherResult: # pragma: no cover 

183 """Decode the given text using the given key. 

184 

185 **Parameters** 

186 - text (str): The text to decode 

187 - key (str): The key to use 

188 

189 **Returns** 

190 - IO: The decoded text 

191 """ 

192 pass 

193 

194 

195registered_cyphers: dict[str, Cypher | CypherWithKey] = {} 

196 

197 

198def register_cypher(cypher: type[Cypher] | type[CypherWithKey]) -> None: 

199 """Class Decorator to register type as cypher. 

200 Must Implment Cypher or CypherWithKey and must have a constructor that takes no args. 

201 

202 **Parameters** 

203 - cypher (type): Type that implements Cypher or CypherWithKey 

204 

205 **Raises** 

206 - ValueError: If cypher does not implement Cypher or CypherWithKey 

207 """ 

208 if not issubclass(cypher, Cypher) and not issubclass(cypher, CypherWithKey): 

209 raise ValueError(f"Invalid cypher Class: {cypher}") 

210 

211 name = cypher.get_name() 

212 

213 if name in registered_cyphers: 

214 logger.warning(f"Duplicate cypher: {name}") 

215 return 

216 

217 registered_cyphers[name] = cypher() 

218 logger.debug(f"Registered cypher: {name}")