Coverage for pki_tools/types/signature_algorithm.py: 79%

1import importlib 

2from abc import ABC 

3from enum import Enum 

4 

5from cryptography.hazmat.primitives import hashes 

6from cryptography.hazmat.primitives._serialization import Encoding 

7from cryptography.hazmat.primitives.asymmetric import ec, padding 

8from cryptography.hazmat.primitives.asymmetric.utils import Prehashed 

9from pydantic import BaseModel, ConfigDict 

10 

11from pki_tools.exceptions import MissingBlockSize 

12from pki_tools.types.crypto_parser import CryptoParser 

13 

14HASHES_MODULE = importlib.import_module( 

15 "cryptography.hazmat.primitives.hashes" 

16) 

17 

18 

19class HashAlgorithmName(Enum): 

20 """ 

21 Enumeration of hash algorithm names. 

22 """ 

23 

24 SHA1 = "SHA1" 

25 SHA512_224 = "SHA512_224" 

26 SHA512_256 = "SHA512_256" 

27 SHA224 = "SHA224" 

28 SHA256 = "SHA256" 

29 SHA384 = "SHA384" 

30 SHA512 = "SHA512" 

31 SHA3_224 = "SHA3_224" 

32 SHA3_256 = "SHA3_256" 

33 SHA3_384 = "SHA3_384" 

34 SHA3_512 = "SHA3_512" 

35 SHAKE128 = "SHAKE128" 

36 SHAKE256 = "SHAKE256" 

37 MD5 = "MD5" 

38 BLAKE2b = "BLAKE2b" 

39 BLAKE2s = "BLAKE2s" 

40 SM3 = "SM3" 

41 

42 

43class HashAlgorithm(CryptoParser): 

44 """ 

45 Represents a hash algorithm. 

46 

47 Attributes: 

48 name: The name of the hash algorithm. 

49 block_size: The block size of the hash algorithm. 

50 """ 

51 

52 name: HashAlgorithmName 

53 block_size: int | None = None 

54 

55 @classmethod 

56 def from_cryptography( 

57 cls: type["HashAlgorithm"], x509_obj: hashes.HashAlgorithm 

58 ) -> "HashAlgorithm": 

59 """ 

60 Constructs a HashAlgorithm object from a cryptography HashAlgorithm 

61 object. 

62 

63 Args: 

64 x509_obj: The cryptography HashAlgorithm object. 

65 

66 Returns: 

67 The constructed HashAlgorithm object. 

68 """ 

69 return cls( 

70 name=getattr(HashAlgorithmName, x509_obj.name.upper()), 

71 block_size=x509_obj.block_size, 

72 _x509_obj=x509_obj, 

73 ) 

74 

75 @property 

76 def der_bytes(self): 

77 return self._to_cryptography().public_bytes(Encoding.DER) 

78 

79 def _to_cryptography(self) -> hashes.HashAlgorithm: 

80 if "SHAKE" in self.name.value: 

81 if self.block_size is None: 

82 raise MissingBlockSize("Please set block_size") 

83 

84 return getattr(HASHES_MODULE, self.name.value)(self.block_size) 

85 

86 if self.name.value == "BLAKE2s": 

87 return getattr(HASHES_MODULE, self.name.value)(32) 

88 elif self.name.value == "BLAKE2b": 

89 return getattr(HASHES_MODULE, self.name.value)(64) 

90 

91 return getattr(HASHES_MODULE, self.name.value)() 

92 

93 def _string_dict(self) -> dict: 

94 return {"algorithm": self.name.value} 

95 

96 

97class Padding(CryptoParser, ABC): 

98 """ 

99 Abstract base class representing padding. 

100 """ 

101 

102 

103class PSSPaddingLenght(Enum): 

104 """ 

105 Enumeration of PSS padding lengths. 

106 """ 

107 

108 _MaxLength = "_MaxLength" 

109 _Auto = "_Auto" 

110 _DigestLength = "_DigestLength" 

111 

112 

113class PSSPadding(Padding): 

114 """ 

115 Represents PSS padding. 

116 

117 Attributes: 

118 mgf: The mask generation function. 

119 length: The length of the padding. 

120 """ 

121 

122 mgf: HashAlgorithm 

123 length: int | PSSPaddingLenght 

124 

125 @classmethod 

126 def from_cryptography( 

127 cls: type["PSSPadding"], crypto_obj: padding.PSS 

128 ) -> "PSSPadding": 

129 """ 

130 Constructs a PSSPadding object from a cryptography PSS object. 

131 

132 Args: 

133 crypto_obj: The cryptography PSS object. 

134 

135 Returns: 

136 The constructed PSSPadding object. 

137 """ 

138 salt_length = crypto_obj._salt_length 

139 if not isinstance(salt_length, int): 

140 salt_length = PSSPaddingLenght[salt_length.__class__.__name__] 

141 

142 return cls( 

143 mgf=crypto_obj._mgf._algorithm, 

144 length=salt_length, 

145 _x509_obj=crypto_obj, 

146 ) 

147 

148 def _to_cryptography(self) -> padding.PSS: 

149 return padding.PSS( 

150 mgf=padding.MGF1(self.mgf._to_cryptography()), 

151 salt_length=getattr(padding, self.length.value), 

152 ) 

153 

154 def _string_dict(self) -> dict: 

155 length = self.length 

156 if not isinstance(length, int): 

157 length = self.length.name 

158 

159 return { 

160 "mgf": self.mgf._string_dict(), 

161 "salt_length": length, 

162 } 

163 

164 

165class PKCS1v15Padding(Padding): 

166 """ 

167 Represents PKCS1v15 padding. 

168 """ 

169 

170 _name: str 

171 

172 @classmethod 

173 def from_cryptography( 

174 cls: type["PKCS1v15Padding"], crypto_obj: padding.PKCS1v15 

175 ) -> "PKCS1v15Padding": 

176 """ 

177 Constructs a PKCS1v15Padding object from a cryptography PKCS1v15 

178 object. 

179 

180 Args: 

181 crypto_obj: The cryptography PKCS1v15 object. 

182 

183 Returns: 

184 The constructed PKCS1v15Padding object. 

185 """ 

186 ret = cls() 

187 ret._name = crypto_obj.name 

188 ret._x509_obj = crypto_obj 

189 return ret 

190 

191 def _to_cryptography(self) -> padding.PKCS1v15: 

192 return padding.PKCS1v15() 

193 

194 def _string_dict(self) -> dict[str, str]: 

195 return {"name": self._name} 

196 

197 

198class ECDSAPadding(Padding): 

199 """ 

200 Represents ECDSA padding. 

201 

202 Attributes: 

203 algorithm: The hash algorithm. 

204 prehashed: Indicates if the data is prehashed. 

205 """ 

206 

207 algorithm: HashAlgorithm 

208 prehashed: bool 

209 

210 @classmethod 

211 def from_cryptography( 

212 cls: type["ECDSAPadding"], crypto_obj: ec.ECDSA 

213 ) -> "ECDSAPadding": 

214 """ 

215 Constructs an ECDSAPadding object from a cryptography ECDSA object. 

216 

217 Args: 

218 crypto_obj: The cryptography ECDSA object. 

219 

220 Returns: 

221 The constructed ECDSAPadding object. 

222 """ 

223 if isinstance(crypto_obj.algorithm, Prehashed): 

224 prehashed = True 

225 algorithm = crypto_obj.algorithm._algorithm 

226 else: 

227 prehashed = False 

228 algorithm = crypto_obj.algorithm 

229 

230 algorithm = HashAlgorithm.from_cryptography(algorithm) 

231 

232 return cls( 

233 algorithm=algorithm, 

234 prehashed=prehashed, 

235 _x509_obj=crypto_obj, 

236 ) 

237 

238 def _to_cryptography(self) -> ec.ECDSA: 

239 alg = self.algorithm._to_cryptography() 

240 if self.prehashed: 

241 alg = Prehashed(alg) 

242 

243 return ec.ECDSA(algorithm=alg) 

244 

245 def _string_dict(self) -> dict: 

246 length = self.length 

247 if not isinstance(length, int): 

248 length = self.length.name 

249 

250 ret = self.algorithm._string_dict() 

251 ret["prehashed"] = str(self.prehashed) 

252 return ret 

253 

254 

255class SignatureAlgorithm(BaseModel): 

256 """ 

257 Represents a signature algorithm. 

258 

259 Attributes: 

260 algorithm: The hash algorithm. 

261 parameters: The parameters of the signature algorithm. 

262 """ 

263 

264 model_config = ConfigDict(arbitrary_types_allowed=True) 

265 

266 algorithm: HashAlgorithm 

267 parameters: PSSPadding | PKCS1v15Padding | ECDSAPadding | None = ( 

268 None 

269 ) 

270 

271 @classmethod 

272 def from_cryptography( 

273 cls: type["SignatureAlgorithm"], 

274 algorithm: hashes.HashAlgorithm, 

275 parameters: padding.PSS | padding.PKCS1v15 | ec.ECDSA = None, 

276 ) -> "SignatureAlgorithm": 

277 """ 

278 Constructs a SignatureAlgorithm object from cryptography objects. 

279 

280 Args: 

281 algorithm: The hash algorithm. 

282 parameters: The parameters. 

283 

284 Returns: 

285 The constructed SignatureAlgorithm object. 

286 """ 

287 algorithm = HashAlgorithm.from_cryptography(algorithm) 

288 

289 if parameters is None: 

290 return cls(algorithm=algorithm) 

291 

292 if isinstance(parameters, padding.PSS): 

293 parameters = PSSPadding.from_cryptography(parameters) 

294 elif isinstance(parameters, padding.PKCS1v15): 

295 parameters = PKCS1v15Padding.from_cryptography(parameters) 

296 elif isinstance(parameters, ec.ECDSA): 

297 parameters = ECDSAPadding.from_cryptography(parameters) 

298 

299 return cls(algorithm=algorithm, parameters=parameters) 

300 

301 def _string_dict(self): 

302 ret = self.algorithm._string_dict() 

303 

304 if self.parameters is not None: 

305 ret["parameters"] = self.parameters._string_dict() 

306 

307 return ret 

308 

309 

310SHA1 = SignatureAlgorithm(algorithm=HashAlgorithm(name=HashAlgorithmName.SHA1)) 

311SHA512_224 = SignatureAlgorithm( 

312 algorithm=HashAlgorithm(name=HashAlgorithmName.SHA512_224) 

313) 

314SHA512_256 = SignatureAlgorithm( 

315 algorithm=HashAlgorithm(name=HashAlgorithmName.SHA512_256) 

316) 

317SHA224 = SignatureAlgorithm( 

318 algorithm=HashAlgorithm(name=HashAlgorithmName.SHA224) 

319) 

320SHA256 = SignatureAlgorithm( 

321 algorithm=HashAlgorithm(name=HashAlgorithmName.SHA256) 

322) 

323SHA384 = SignatureAlgorithm( 

324 algorithm=HashAlgorithm(name=HashAlgorithmName.SHA384) 

325) 

326SHA512 = SignatureAlgorithm( 

327 algorithm=HashAlgorithm(name=HashAlgorithmName.SHA512) 

328) 

329SHA3_224 = SignatureAlgorithm( 

330 algorithm=HashAlgorithm(name=HashAlgorithmName.SHA3_224) 

331) 

332SHA3_256 = SignatureAlgorithm( 

333 algorithm=HashAlgorithm(name=HashAlgorithmName.SHA3_256) 

334) 

335SHA3_384 = SignatureAlgorithm( 

336 algorithm=HashAlgorithm(name=HashAlgorithmName.SHA3_384) 

337) 

338SHA3_512 = SignatureAlgorithm( 

339 algorithm=HashAlgorithm(name=HashAlgorithmName.SHA3_512) 

340) 

341SHAKE128 = SignatureAlgorithm( 

342 algorithm=HashAlgorithm(name=HashAlgorithmName.SHAKE128, block_size=64) 

343) 

344SHAKE256 = SignatureAlgorithm( 

345 algorithm=HashAlgorithm(name=HashAlgorithmName.SHAKE256, block_size=64) 

346) 

347MD5 = SignatureAlgorithm(algorithm=HashAlgorithm(name=HashAlgorithmName.MD5)) 

348BLAKE2b = SignatureAlgorithm( 

349 algorithm=HashAlgorithm(name=HashAlgorithmName.BLAKE2b) 

350) 

351BLAKE2s = SignatureAlgorithm( 

352 algorithm=HashAlgorithm(name=HashAlgorithmName.BLAKE2s) 

353) 

354SM3 = SignatureAlgorithm(algorithm=HashAlgorithm(name=HashAlgorithmName.SM3))