66
77from __future__ import annotations
88
9+ import abc
910import argparse
1011import sys
1112from hashlib import sha256 , sha512
12- from typing import BinaryIO
13+ from typing import BinaryIO , Type
1314
1415from cryptography .exceptions import InvalidSignature
15- from cryptography .hazmat .primitives import hashes
16- from cryptography .hazmat .primitives import serialization
17- from cryptography .hazmat .primitives .asymmetric import ec
18- from cryptography .hazmat .primitives .asymmetric import ed25519
16+ from cryptography .hazmat .primitives import hashes , serialization
17+ from cryptography .hazmat .primitives .asymmetric import ec , ed25519
1918from cryptography .hazmat .primitives .serialization import load_pem_private_key
2019
2120
@@ -62,44 +61,53 @@ def generate_legal_key_for_ed25519():
6261 return key
6362
6463
65- class EllipticCurveKeysGenerator :
66- """Generate private and public keys for Elliptic Curve cryptography."""
64+ class KeysGeneratorBase (abc .ABC ):
6765
6866 def __init__ (self , infile : BinaryIO | None = None ) -> None :
6967 """
7068 :param infile: A file-like object to read the private key.
7169 """
7270 if infile is None :
73- self .private_key = generate_legal_key_for_elliptic_curve ()
71+ self .private_key = self . _generate_private_key ()
7472 else :
7573 self .private_key = load_pem_private_key (infile .read (), password = None )
7674 self .public_key = self .private_key .public_key ()
7775
7876 @property
77+ @abc .abstractmethod
7978 def private_key_pem (self ) -> bytes :
80- return self .private_key .private_bytes (
81- encoding = serialization .Encoding .PEM ,
82- format = serialization .PrivateFormat .PKCS8 ,
83- encryption_algorithm = serialization .NoEncryption (),
84- )
79+ pass
80+
81+ @property
82+ @abc .abstractmethod
83+ def public_key_pem (self ) -> bytes :
84+ pass
85+
86+ @abc .abstractmethod
87+ def _generate_private_key (self ):
88+ pass
89+
90+ @staticmethod
91+ @abc .abstractmethod
92+ def sign_message (private_key , message : bytes ) -> bytes :
93+ pass
94+
95+ @staticmethod
96+ @abc .abstractmethod
97+ def verify_signature (
98+ public_key , message : bytes , signature : bytes
99+ ) -> bool :
100+ pass
85101
86102 def write_private_key_pem (self , outfile : BinaryIO ) -> bytes :
87103 """
88104 Write private key pem to file and return it.
89105
90106 :param outfile: A file-like object to write the private key.
91107 """
92- if outfile is not None :
93- outfile .write (self .private_key_pem )
108+ outfile .write (self .private_key_pem )
94109 return self .private_key_pem
95110
96- @property
97- def public_key_pem (self ) -> bytes :
98- return self .public_key .public_bytes (
99- encoding = serialization .Encoding .PEM ,
100- format = serialization .PublicFormat .SubjectPublicKeyInfo ,
101- )
102-
103111 def write_public_key_pem (self , outfile : BinaryIO ) -> bytes :
104112 """
105113 Write public key pem to file and return it.
@@ -109,31 +117,48 @@ def write_public_key_pem(self, outfile: BinaryIO) -> bytes:
109117 outfile .write (self .public_key_pem )
110118 return self .public_key_pem
111119
120+
121+ class EllipticCurveKeysGenerator (KeysGeneratorBase ):
122+ """Generate private and public keys for Elliptic Curve cryptography."""
123+
124+ def _generate_private_key (self ):
125+ return generate_legal_key_for_elliptic_curve ()
126+
127+ @property
128+ def private_key_pem (self ) -> bytes :
129+ return self .private_key .private_bytes (
130+ encoding = serialization .Encoding .PEM ,
131+ format = serialization .PrivateFormat .PKCS8 ,
132+ encryption_algorithm = serialization .NoEncryption (),
133+ )
134+
135+ @property
136+ def public_key_pem (self ) -> bytes :
137+ return self .public_key .public_bytes (
138+ encoding = serialization .Encoding .PEM ,
139+ format = serialization .PublicFormat .SubjectPublicKeyInfo ,
140+ )
141+
112142 @staticmethod
113- def verify_signature (public_key , message : bytes , signature : bytes ) -> bool :
143+ def verify_signature (
144+ public_key : ec .EllipticCurvePublicKey , message : bytes , signature : bytes
145+ ) -> bool :
114146 try :
115147 public_key .verify (signature , message , ec .ECDSA (hashes .SHA256 ()))
116148 return True
117149 except InvalidSignature :
118150 return False
119151
120152 @staticmethod
121- def sign_message (private_key , message : bytes ) -> bytes :
153+ def sign_message (private_key : ec . EllipticCurvePrivateKey , message : bytes ) -> bytes :
122154 return private_key .sign (message , ec .ECDSA (hashes .SHA256 ()))
123155
124156
125- class Ed25519KeysGenerator :
157+ class Ed25519KeysGenerator ( KeysGeneratorBase ) :
126158 """Generate private and public keys for ED25519 cryptography."""
127159
128- def __init__ (self , infile : BinaryIO | None = None ) -> None :
129- """
130- :param infile: A file-like object to read the private key.
131- """
132- if infile is None :
133- self .private_key : ed25519 .Ed25519PrivateKey = generate_legal_key_for_ed25519 ()
134- else :
135- self .private_key = load_pem_private_key (infile .read (), password = None ) # type: ignore[assignment]
136- self .public_key : ed25519 .Ed25519PublicKey = self .private_key .public_key ()
160+ def _generate_private_key (self ):
161+ return generate_legal_key_for_ed25519 ()
137162
138163 @property
139164 def private_key_pem (self ) -> bytes :
@@ -143,50 +168,40 @@ def private_key_pem(self) -> bytes:
143168 encryption_algorithm = serialization .NoEncryption ()
144169 )
145170
146- def write_private_key_pem (self , outfile : BinaryIO ) -> bytes :
147- """
148- Write private key pem to file and return it.
149-
150- :param outfile: A file-like object to write the private key.
151- """
152- outfile .write (self .private_key_pem )
153- return self .private_key_pem
154-
155171 @property
156172 def public_key_pem (self ) -> bytes :
157173 return self .public_key .public_bytes (
158174 encoding = serialization .Encoding .PEM ,
159175 format = serialization .PublicFormat .SubjectPublicKeyInfo
160176 )
161177
162- def write_public_key_pem (self , outfile : BinaryIO ) -> bytes :
163- """
164- Write public key pem to file and return it.
165-
166- :param outfile: A file-like object to write the public key.
167- """
168- if outfile is not None :
169- outfile .write (self .public_key_pem )
170- return self .public_key_pem
171-
172178 @staticmethod
173- def verify_signature (public_key : ed25519 .Ed25519PublicKey , message : bytes , signature : bytes ) -> bool :
179+ def verify_signature (
180+ public_key : ed25519 .Ed25519PublicKey , message : bytes , signature : bytes
181+ ) -> bool :
174182 try :
175183 public_key .verify (signature , message )
176184 return True
177185 except InvalidSignature :
178186 return False
179187
180188 @staticmethod
181- def sign_message (private_key , message : bytes ) -> bytes :
189+ def sign_message (private_key : ed25519 . Ed25519PrivateKey , message : bytes ) -> bytes :
182190 return private_key .sign (message )
183191
184192
193+ ALGORITHMS : dict [str , Type [KeysGeneratorBase ]] = {
194+ "ed25519" : Ed25519KeysGenerator ,
195+ "ec" : EllipticCurveKeysGenerator ,
196+ }
197+
198+
185199def main (argv = None ) -> int :
186200 parser = argparse .ArgumentParser (
187201 description = 'Generate PEM file.' ,
188202 formatter_class = argparse .RawDescriptionHelpFormatter ,
189- allow_abbrev = False )
203+ allow_abbrev = False
204+ )
190205
191206 priv_pub_group = parser .add_mutually_exclusive_group (required = True )
192207 priv_pub_group .add_argument ('--private' , required = False , action = 'store_true' ,
@@ -207,19 +222,15 @@ def main(argv=None) -> int:
207222
208223 args = parser .parse_args (argv )
209224
210- if args .algorithm == 'ed25519' :
211- ed25519_generator = Ed25519KeysGenerator (args .infile )
212- if args .private :
213- ed25519_generator .write_private_key_pem (args .out )
214- if args .public :
215- ed25519_generator .write_public_key_pem (args .out )
216- else :
217- ec_generator = EllipticCurveKeysGenerator (args .infile )
218- if args .private :
219- ec_generator .write_private_key_pem (args .out )
220- elif args .public :
221- ec_generator .write_public_key_pem (args .out )
225+ try :
226+ generator = ALGORITHMS [args .algorithm ](args .infile )
227+ except KeyError :
228+ sys .exit (f'Unknown algorithm { args .algorithm } .' )
222229
230+ if args .private :
231+ generator .write_private_key_pem (args .out )
232+ if args .public :
233+ generator .write_public_key_pem (args .out )
223234 return 0
224235
225236
0 commit comments