pyjod/pyjod/appdata.py

import json
from base64 import b64encode, b64decode
from pathlib import Path

import xdgappdirs
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import serialization
from jwcrypto import jwk
from pathvalidate import sanitize_filename

data_dir = xdgappdirs.user_data_dir('pyjod', as_path=True)


class RSAPrivateKey:
    def __init__(self, key):
        self.key = key

    @classmethod
    def generate(cls):
        key = rsa.generate_private_key(65537, 2048)
        return cls(key)

    @classmethod
    def from_pem(cls, data):
        key = serialization.load_pem_private_key(data, None)
        return cls(key)

    @property
    def jwk(self):
        return jwk.JWK.from_pyca(self.key)

    @property
    def pubkey_b64(self):
        pubkey = self.key.public_key()
        pubkey_bytes = pubkey.public_bytes(
            serialization.Encoding.DER,
            serialization.PublicFormat.SubjectPublicKeyInfo
        )
        return b64encode(pubkey_bytes).decode('utf-8')

    def to_pem(self):
        return self.key.private_bytes(
            serialization.Encoding.PEM,
            serialization.PrivateFormat.TraditionalOpenSSL,
            serialization.NoEncryption()
        )


class RSAPublicKey:
    def __init__(self, key):
        self.key = key

    @classmethod
    def from_b64(cls, data):
        key_bytes = b64decode(data)
        key = serialization.load_der_public_key(key_bytes)
        return cls(key)

    @classmethod
    def from_pem(cls, data):
        key = serialization.load_pem_public_key(data)
        return cls(key)

    @property
    def jwk(self):
        return jwk.JWK.from_pyca(self.key)

    def to_pem(self):
        return self.key.public_bytes(
            serialization.Encoding.PEM,
            serialization.PublicFormat.SubjectPublicKeyInfo
        )


class AppData:
    def __init__(self, profile_name):
        self.profile_dir = data_dir / sanitize_filename(profile_name)
        self.profile_dir.mkdir(exist_ok=True, parents=True)

        self.values_file = self.profile_dir / "values.json"

    @property
    def app_privkey(self):
        key_file = self.profile_dir / 'app_privkey.pem'
        if key_file.is_file():
            with key_file.open('rb') as f:
                key_bytes = f.read()
            return RSAPrivateKey.from_pem(key_bytes)
        else:
            key = RSAPrivateKey.generate()
            with key_file.open('wb') as f:
                f.write(key.to_pem())
            return key

    @property
    def serv_pubkey(self):
        key_file = self.profile_dir / 'serv_pubkey.pem'
        if key_file.is_file():
            with key_file.open('rb') as f:
                key_bytes = f.read()
            return RSAPublicKey.from_pem(key_bytes)
        else:
            return None

    @serv_pubkey.setter
    def serv_pubkey(self, key_b64):
        key = RSAPublicKey.from_b64(key_b64)
        key_file = self.profile_dir / 'serv_pubkey.pem'
        with key_file.open('wb') as f:
            f.write(key.to_pem())

    def __getitem__(self, key):
        if self.values_file.is_file():
            with self.values_file.open() as f:
                values = json.load(f)
            if key in values:
                return values[key]
            return None

    def __setitem__(self, key, value):
        if self.values_file.is_file():
            with self.values_file.open() as f:
                values = json.load(f)
        else:
            values = {}
        values[key] = value
        with self.values_file.open('w') as f:
            json.dump(values, f)