communex.compat.storage

  1"""
  2Data storage compatible with the *classic* `commune` library.
  3"""
  4
  5# TODO: encryption
  6
  7import base64
  8import hashlib
  9import json
 10import os.path
 11import time
 12from typing import Any
 13
 14from nacl.secret import SecretBox
 15from nacl.utils import random
 16
 17from communex.errors import PasswordNotProvidedError
 18from communex.util import ensure_parent_dir_exists
 19
 20# from cryptography.fernet import Fernet
 21# from cryptography.hazmat.primitives import hashes
 22# from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 23
 24
 25COMMUNE_HOME = "~/.commune"
 26"""
 27COMMUNE_HOME
 28
 29    classic commune data storage home directory.
 30"""
 31
 32
 33def _derive_key(password: str):
 34    # Derive a 256-bit key from the password using Blake2b
 35    key = hashlib.blake2b(password.encode(), digest_size=32).digest()
 36    return key
 37
 38
 39def _encrypt_data(password: str, data: Any) -> str:
 40    key = _derive_key(password)
 41    box = SecretBox(key)
 42    nonce = random(SecretBox.NONCE_SIZE)
 43    raw = json.dumps(data).encode()
 44    ciphertext = box.encrypt(raw, nonce).ciphertext
 45    encrypted = nonce + ciphertext
 46    decoded_data = base64.b64encode(encrypted).decode()
 47    return decoded_data
 48
 49
 50def _decrypt_data(password: str, data: str) -> Any:
 51    key = _derive_key(password)
 52    box = SecretBox(key)
 53    encrypted = base64.b64decode(data.encode())
 54    nonce = encrypted[: SecretBox.NONCE_SIZE]
 55    ciphertext = encrypted[SecretBox.NONCE_SIZE :]
 56    raw = box.decrypt(ciphertext, nonce)
 57    return json.loads(raw.decode())
 58
 59
 60def classic_load(
 61    path: str, mode: str = "json", password: str | None = None
 62) -> Any:
 63    """
 64    Load data from commune data storage.
 65
 66    Args:
 67        path: Data storage file path.
 68        mode: Data storage mode.
 69
 70    Returns:
 71        Data loaded from the data storage.
 72
 73    Todo:
 74        * Other serialization modes support. Only json mode is supported now.
 75
 76    Raises:
 77        NotImplementedError: See Todo.
 78        AssertionError: Raised when the data is not in the classic format.
 79    """
 80    if mode != "json":
 81        raise NotImplementedError(
 82            "Our commune data storage only supports json mode"
 83        )
 84
 85    full_path = os.path.expanduser(os.path.join(COMMUNE_HOME, path))
 86    with open(full_path, "r") as file:
 87        body = json.load(file)
 88
 89    if body["encrypted"] and password is None:
 90        raise PasswordNotProvidedError(
 91            "Data is encrypted but no password provided"
 92        )
 93    if body["encrypted"] and password is not None:
 94        content = _decrypt_data(password, body["data"])
 95    else:
 96        content = body["data"]
 97
 98    assert isinstance(body, dict)
 99    assert isinstance(body["timestamp"], int)
100    assert isinstance(content, (dict, list, tuple, set, float, str, int))
101    return content  # type: ignore
102
103
104def classic_put(
105    path: str, value: Any, mode: str = "json", password: str | None = None
106):
107    """
108    Put data into commune data storage.
109
110    Args:
111        path: Data storage path.
112        value: Data to store.
113        mode: Data storage mode.
114        encrypt: Whether to encrypt the data.
115
116    Todo:
117        * Other serialization modes support. Only json mode is supported now.
118
119    Raises:
120        NotImplementedError: See Todo.
121        TypeError: Raised when value is not a valid type.
122        FileExistsError: Raised when the file already exists.
123    """
124    if mode != "json":
125        raise NotImplementedError(
126            "Our commune data storage only supports json mode"
127        )
128    if not isinstance(value, (dict, list, tuple, set, float, str, int)):
129        raise TypeError(
130            f"Invalid type for commune data storage value: {type(value)}"
131        )
132
133    timestamp = int(time.time())
134
135    full_path = os.path.expanduser(os.path.join(COMMUNE_HOME, path))
136
137    if os.path.exists(full_path):
138        raise FileExistsError(
139            f"Commune data storage file already exists: {full_path}"
140        )
141
142    ensure_parent_dir_exists(full_path)
143
144    if password:
145        value = _encrypt_data(password, value)
146        encrypted = True
147    else:
148        encrypted = False
149
150    with open(full_path, "w") as file:
151        json.dump(
152            {"data": value, "encrypted": encrypted, "timestamp": timestamp},
153            file,
154            indent=4,
155        )