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