activity_tools.crypto

 1import os
 2
 3from cryptography.hazmat.primitives import serialization
 4from cryptography.hazmat.primitives.asymmetric import rsa
 5from cryptography.hazmat.backends import default_backend
 6
 7class RSAKey():
 8    """
 9    This class holds private and public RSA keypairs with file
10    load and save functionality. This is an utility class that
11    you are free to ignore if you prefer to use your own
12    cryptographic functions.
13
14    ```python
15    Example:
16        key = RSAKey("/my/path/key.pem")
17        print(key.public_key)
18        signed = key.sign(text, padding.PKCS1v15(), hashes.SHA256())
19    ```
20    """
21
22    def __init__(self, path) -> None:
23        """
24        This creates/loads/writes a private key to the specified path.
25        You can access both the private and public keys from the members
26        `.private_key` and `.public_key`.
27        """
28        self.private_key_path = path
29
30        if os.path.exists(path):
31            key = self.load_private_key(path)
32        else:
33            key = self.generate_private_key()
34            self.save_private_key(path)
35            print(f"Saved newly generated keys to {path}")
36
37        self.key: rsa.RSAPrivateKey = key
38        """ The RSAPrivateKey object from cryptography """
39
40        self.private_key: bytes = self.key.private_bytes(
41            serialization.Encoding.PEM,
42            serialization.PrivateFormat.PKCS8,
43            serialization.NoEncryption())
44        """ Contains the private key """
45
46        self.public_key: bytes = self.key.public_key().public_bytes(
47            serialization.Encoding.PEM,
48            serialization.PublicFormat.SubjectPublicKeyInfo
49        )
50        """ Contains the public key """
51
52    def generate_private_key(self) -> rsa.RSAPrivateKey:
53        """
54        Generate a RSA private key and returns it. The constructor calles this
55        to generate a private key.
56        """
57        return rsa.generate_private_key(
58            public_exponent=65537, key_size=2048, backend=default_backend()
59        )
60
61    def save_private_key(self, path: str) -> None:
62        """
63        Export the private key to bytes and write them to the specified
64        filepath. The constructor uses this to write the keys to storate.
65        """
66        pem = self.key.private_bytes(
67            encoding=serialization.Encoding.PEM,
68            format=serialization.PrivateFormat.PKCS8,
69            encryption_algorithm=serialization.NoEncryption()
70        )
71        
72        with open(path, 'wb') as pem_out:
73            pem_out.write(pem)
74
75    def load_private_key(self, path: str) -> rsa.RSAPrivateKey:
76        """
77        Import and load a private key from a file.
78        """
79        with open(path, "rb") as key_file:
80            key = serialization.load_pem_private_key(
81                key_file.read(),
82                password=None,
83                backend=default_backend()
84        )
85        return key