Generating key stores for database

gen_key_store.py

import bcrypt
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import utils

class settings(object):
    SECURE_STORE = {
        'PRIVATE_KEY_FILE': None,
        'PUBLIC_KEY': None,
        'PUBLIC_KEY_FILE': None,
        'PUBLIC_EXPONENT': 65337,
        'KEY_SIZE': 2048,
        'HASH_SIZE': 256,
        'BLOCK_SIZE': 190
    }

_BACKEND = default_backend()
if settings.SECURE_STORE['HASH_SIZE'] == 256:
    _ALGORITHM = hashes.SHA256()
else:
    raise ValueError('Unrecognized hash size in settings (expected one of (256,))')
_PAD = padding.OAEP(
    mgf=padding.MGF1(algorithm=_ALGORITHM),
    algorithm=_ALGORITHM,
    label=None
)

def gen_salt():
    return bcrypt.gensalt()

def get_private_key_bytes(private_key, password=None, salt=None):
    encryption = serialization.NoEncryption()
    if password:
        if salt:
            encryption = serialization.BestAvailableEncryption(salt + password)
        else:
            encryption = serialization.BestAvailableEncryption(password)

    return private_key.private_bytes(
        serialization.Encoding.PEM,
        serialization.PrivateFormat.PKCS8,
        encryption
    )

def get_public_key_bytes(public_key):
    return public_key.public_bytes(
        serialization.Encoding.PEM,
        serialization.PublicFormat.PKCS1
    )

def gen_keys(password=None, salt=None):
    private_key = rsa.generate_private_key(
        public_exponent=settings.SECURE_STORE['PUBLIC_EXPONENT'],
        key_size=settings.SECURE_STORE['KEY_SIZE'],
        backend=_BACKEND
    )

    private_key_bytes = get_private_key_bytes(private_key, password, salt)

    public_key = private_key.public_key()

    public_key_bytes = get_public_key_bytes(public_key)

    return private_key_bytes, public_key_bytes

def load_keys(prk, puk, password=None, salt=None):
    pwd = None
    if password:
        if salt:
            pwd = salt + password
        else:
            pwd = password

    private_key = serialization.load_pem_private_key(
        prk,
        password=pwd,
        backend=default_backend()
    )

    public_key = serialization.load_pem_public_key(
        puk,
        backend=default_backend()
    )

    return private_key, public_key

def aencrypt(message, public_key):
    enc = b''
    for i in range(0, len(message), settings.SECURE_STORE['BLOCK_SIZE']):
        enc += public_key.encrypt(message[i:i + settings.SECURE_STORE['BLOCK_SIZE']], _PAD)
    return enc

def adecrypt(encrypted, private_key):
    dec = b''
    for i in range(0, len(encrypted), settings.SECURE_STORE['HASH_SIZE']):
        dec += private_key.decrypt(encrypted[i:i + settings.SECURE_STORE['HASH_SIZE']], _PAD)
    return dec

def encrypt_key_store(private_key, public_key, password=None, salt=None, server_public_key=None):
    return (
        salt,
        bcrypt.hashpw(password, salt),
        get_public_key_bytes(public_key),
        aencrypt(get_private_key_bytes(private_key, password, salt), server_public_key)
        if server_public_key
        else get_private_key_bytes(private_key, password, salt)
    )

def decrypt_key_store(password, salt, hashed, public_key_bytes, private_key_bytes, server_private_key=None):
    if password:
        if not bcrypt.checkpw(password, hashed):
            raise ValueError('Invalid password')
    if server_private_key:
        prk = adecrypt(private_key_bytes, server_private_key)
    else:
        prk = private_key_bytes
    return load_keys(prk, public_key_bytes, password, salt)

def gen_key_store(password, server_public_key=None):
    salt = gen_salt()
    prk, puk = load_keys(*gen_keys(password, salt), password, salt)
    return encrypt_key_store(prk, puk, password, salt, server_public_key)

sprk, spuk = load_keys(*gen_keys())
a = gen_key_store(b'asdf', spuk)  # This is what can be safely stored in the database
b = decrypt_key_store(b'asdf', *a, sprk)

#### TODO:
# Load private / public keys for server from file
# Comments explaining everything
# Actual test script
# Eh

For example, calling the following:
gen_key_store(b'This is my abnormally long password because I care about security hey!23456790%$234./p[][][{}{}[][]@$#(@*#%')
generates the following output, which is in the form salt, hashed, public_key_bytes, private_key_encrypted_by_password_that_can_be_encrypted_by_server_bytes:

b'$2b$12$P4Tow9aXvrY228Crga90aO'

b'$2b$12$P4Tow9aXvrY228Crga90aOJdIk7nlfIUujSrru/HLw91enEBLOYjG'

b'-----BEGIN RSA PUBLIC KEY-----\nMIIBCgKCAQEAtRKHm3Oz2/HiaZY0TsxdWLLmcGz/BfZWZw15U2NX5qS7QbiiImww\ncKzltxd49D2c5B3W0yZNCbK9CKp8idK5usxtj4+pqq7oMX53nuZnd2cw+Cb27hdD\nARkdNGa6HFBa1qtgvuiJ+WfIXdwovrgau813K/KdQN6e7xZ0xe0ol/QRu97mmmiK\nxGoZIeE+yDJ7oIQtk0CFu3+UH9D27JCavr5SrmZLfevskfbVdJ5xtcxikfo7uHJx\nJT/3i5fbDzm3Z0eJEC3pYf9GYTo+fqcS6PB1BWXkBbF7Gy3guy9spDRb6vPyolVh\nSWlyPFfN0Y2O3HVRhzLJ/qPH5LJREru7QQIDAP85\n-----END RSA PUBLIC KEY-----\n'

b'-----BEGIN ENCRYPTED PRIVATE KEY-----\nMIIFHzBJBgkqhkiG9w0BBQ0wPDAbBgkqhkiG9w0BBQwwDgQIZPHT9Qt7iG8CAggA\nMB0GCWCGSAFlAwQBKgQQJBo64R2/9kM2ZsCCTd2wXASCBNA+G43jQE9yDcGnwWmZ\n9kuPRb590M+o3M+b33cX1Fd3FY8CDtmp7LpS0JpUy1YZXkiZmsdg3IjKk+GPT2UB\nss+HahDz4LAoU1vRgCw2Be6nFYQaLQxwkA0M43hhbNeLMyIh0vOVG3w9yzbcZjfe\nC3YLNB410D0ObOjmg9tnIa7VR1PyC5L1Dev0TjgwfKy3vaAsbMTuzr7Yu5Zi6Rtq\nsowDor1AK0eqsGeVafgk23RRvI5wLESK7S9ZYgi/EbqqQxOG5NQcBOMShM0NW8hP\nxzE8gEi1zkE1Ej5ZxWi/FSQ+OGGHPpXvs7Q+Dr0mEwKt32uvlV++zM7/OQlhpZTw\nMBJe6bmKUOkn7ZfzqCYbbmOhfnEqw4Sc5XZ2J67+DyspYsXWywiZpAaAQEkqrtx6\nbU5tm0i79mgZNnm9XftDpNz4Xo+KTt/huRJgPYX2bPfuVTyvKcoyNfUyHanZ1f8F\nnzOnl3vd27R0/FD44OYD4Ja9U5drwkygum7PUvbsvkkT/nWUFN7cEXOROjU3dNlP\nLx+TlF6mmJD0yHtF1VbOquNEPUrEg34GZm0/K6XBBXcHBu5cXaqiRekVrs2xrOvG\nVXZFWj83Vnfz1QfZS7uEnOz0NRLo3mkuO9FGCCZQOPQ5NDaFt9rnRYF5dXZpAwkM\nN2o231l6zuReQPhdjrJNhCWCFnivJlsJDZwwYAbmHYCo5N6giK6/8EL/AqCdYEsR\nOIrz8ROvwcmX1QlKm2kno6V3L4AyPyBvDDadRkEl/RS9fLP/q+Wqic+ODO3qFczN\nui+YawfBB+0wJAp6qXTsdBWyKknLfhGo1K4DVIJeAdMwMALuFC+SY9+cLTTpfY+F\nRIJde6HnrJHxGALzonLt1QD1LGtNzbw38YgdEHEcgRXprzymRsUymJ+tFL9LB+5s\n3JR/nmsXsTmkWVGr7wFAEGnON+xxktxkoInt8pmgzJHWePjVwXN9fID+nN99Pvyz\njoxtsLFFPmDC0vTx0ZIdDPdbD5yAbkMBvdx8bGWlq+JM9iOn9XIFImBBi8IXepUp\n/7gQqnC9zpHvevfc+bmpDlB2DnJOvLKAr2zNFDAyd/OKrG3ClqQEqUc40E90bkV1\nQikn1htPoJehrlhZzoFcp/JFAwr2f0XdYkSzH8kyG03J2GB+PlJQYt8iozIIneEE\nqvRsdGkGCnk2UPd+H3ipmEGbvsY1yOrZqEoV1XLy0Qbt7mrRmkcGMjcHqTp7muxY\nR099AxmQDiuXnVoYJq7NiLNNKfM8oCVMsZYCiMA6wu8waESsX6vXuRVpuZRyR7rr\nZmJoigfLbp1cKET2b30XBY1om4Dwu5TrlOKCIwH4WTpoSs04tvrX2OLbXesgYUo7\nsf1B8QAnQhyriU/hVRjWU0UuPgMDEAfuKzAdpxjEKC+1XiOdhvKnNjipQmO3HoH6\nDTwWcn6yd7BOcDGuI6JzG73wGZddHxN63qs6DwiQPAuhBBbtPm6rcOZUR6CMRIiz\n1meB0QoyDvBB+6KvwgeoAvbirfMHNvh9PBJZ4YIE4EzYdHHLKC3wJjimpPIFClvE\nlyECK80BjwOKqqfrq48KLcmo+hzn+FF5G6mFQqujZd0m1u5rqZZue/G2vEOqSzuG\nKo+d4qZ9Fd7F1grLx20egSzZTg==\n-----END ENCRYPTED PRIVATE KEY-----\n'```