aboutsummaryrefslogtreecommitdiff
path: root/.venv/lib/python3.12/site-packages/core/providers/crypto/nacl.py
blob: 63232565dccab0f0b6e66f7e2af269dd38317656 (about) (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
import base64
import logging
import os
import string
from datetime import datetime, timezone
from typing import Optional, Tuple

import jwt
import nacl.encoding
import nacl.exceptions
import nacl.pwhash
import nacl.signing
from nacl.exceptions import BadSignatureError
from nacl.pwhash import argon2i

from core.base import CryptoConfig, CryptoProvider

DEFAULT_NACL_SECRET_KEY = "wNFbczH3QhUVcPALwtWZCPi0lrDlGV3P1DPRVEQCPbM"  # Replace or load from env or secrets manager


def encode_bytes_readable(random_bytes: bytes, chars: str) -> str:
    """Convert random bytes to a readable string using the given character
    set."""
    # Each byte gives us 8 bits of randomness
    # We use modulo to map each byte to our character set
    result = []
    for byte in random_bytes:
        # Use modulo to map the byte (0-255) to our character set length
        idx = byte % len(chars)
        result.append(chars[idx])
    return "".join(result)


class NaClCryptoConfig(CryptoConfig):
    provider: str = "nacl"
    # Interactive parameters for password ops (fast)
    ops_limit: int = argon2i.OPSLIMIT_MIN
    mem_limit: int = argon2i.MEMLIMIT_MIN
    # Sensitive parameters for API key generation (slow but more secure)
    api_ops_limit: int = argon2i.OPSLIMIT_INTERACTIVE
    api_mem_limit: int = argon2i.MEMLIMIT_INTERACTIVE
    api_key_bytes: int = 32
    secret_key: Optional[str] = None


class NaClCryptoProvider(CryptoProvider):
    def __init__(self, config: NaClCryptoConfig):
        if not isinstance(config, NaClCryptoConfig):
            raise ValueError(
                "NaClCryptoProvider must be initialized with a NaClCryptoConfig"
            )
        super().__init__(config)
        self.config: NaClCryptoConfig = config
        logging.info("Initializing NaClCryptoProvider")

        # Securely load the secret key for JWT
        # Priority: config.secret_key > environment variable > default
        self.secret_key = (
            config.secret_key
            or os.getenv("R2R_SECRET_KEY")
            or DEFAULT_NACL_SECRET_KEY
        )

    def get_password_hash(self, password: str) -> str:
        password_bytes = password.encode("utf-8")
        hashed = nacl.pwhash.argon2i.str(
            password_bytes,
            opslimit=self.config.ops_limit,
            memlimit=self.config.mem_limit,
        )
        return base64.b64encode(hashed).decode("utf-8")

    def verify_password(
        self, plain_password: str, hashed_password: str
    ) -> bool:
        try:
            stored_hash = base64.b64decode(hashed_password.encode("utf-8"))
            nacl.pwhash.verify(stored_hash, plain_password.encode("utf-8"))
            return True
        except nacl.exceptions.InvalidkeyError:
            return False

    def generate_verification_code(self, length: int = 32) -> str:
        random_bytes = nacl.utils.random(length)
        return base64.urlsafe_b64encode(random_bytes)[:length].decode("utf-8")

    def generate_api_key(self) -> Tuple[str, str]:
        # Define our character set (excluding ambiguous characters)
        chars = string.ascii_letters.replace("l", "").replace("I", "").replace(
            "O", ""
        ) + string.digits.replace("0", "").replace("1", "")

        # Generate a unique key_id
        key_id_bytes = nacl.utils.random(16)  # 16 random bytes
        key_id = f"pk_{encode_bytes_readable(key_id_bytes, chars)}"

        # Generate a high-entropy API key
        raw_api_key = f"sk_{encode_bytes_readable(nacl.utils.random(self.config.api_key_bytes), chars)}"

        # The caller will store the hashed version in the database
        return key_id, raw_api_key

    def hash_api_key(self, raw_api_key: str) -> str:
        hashed = nacl.pwhash.argon2i.str(
            raw_api_key.encode("utf-8"),
            opslimit=self.config.api_ops_limit,
            memlimit=self.config.api_mem_limit,
        )
        return base64.b64encode(hashed).decode("utf-8")

    def verify_api_key(self, raw_api_key: str, hashed_key: str) -> bool:
        try:
            stored_hash = base64.b64decode(hashed_key.encode("utf-8"))
            nacl.pwhash.verify(stored_hash, raw_api_key.encode("utf-8"))
            return True
        except nacl.exceptions.InvalidkeyError:
            return False

    def sign_request(self, private_key: str, data: str) -> str:
        try:
            key_bytes = base64.b64decode(private_key)
            signing_key = nacl.signing.SigningKey(key_bytes)
            signature = signing_key.sign(data.encode())
            return base64.b64encode(signature.signature).decode()
        except Exception as e:
            raise ValueError(
                f"Invalid private key or signing error: {str(e)}"
            ) from e

    def verify_request_signature(
        self, public_key: str, signature: str, data: str
    ) -> bool:
        try:
            key_bytes = base64.b64decode(public_key)
            verify_key = nacl.signing.VerifyKey(key_bytes)
            signature_bytes = base64.b64decode(signature)
            verify_key.verify(data.encode(), signature_bytes)
            return True
        except (BadSignatureError, ValueError):
            return False

    def generate_secure_token(self, data: dict, expiry: datetime) -> str:
        """Generate a secure token using JWT with HS256.

        The secret_key is used for symmetrical signing.
        """
        now = datetime.now(timezone.utc)
        to_encode = {
            **data,
            "exp": expiry.timestamp(),
            "iat": now.timestamp(),
            "nbf": now.timestamp(),
            "jti": base64.urlsafe_b64encode(nacl.utils.random(16)).decode(),
            "nonce": base64.urlsafe_b64encode(nacl.utils.random(16)).decode(),
        }

        return jwt.encode(to_encode, self.secret_key, algorithm="HS256")

    def verify_secure_token(self, token: str) -> Optional[dict]:
        """Verify a secure token using the shared secret_key and JWT."""
        try:
            payload = jwt.decode(token, self.secret_key, algorithms=["HS256"])
            exp = payload.get("exp")
            if exp is None or datetime.fromtimestamp(
                exp, tz=timezone.utc
            ) < datetime.now(timezone.utc):
                return None
            return payload
        except (jwt.ExpiredSignatureError, jwt.InvalidTokenError):
            return None

    def generate_signing_keypair(self) -> Tuple[str, str, str]:
        signing_key = nacl.signing.SigningKey.generate()
        private_key_b64 = base64.b64encode(signing_key.encode()).decode()
        public_key_b64 = base64.b64encode(
            signing_key.verify_key.encode()
        ).decode()
        # Generate a unique key_id
        key_id_bytes = nacl.utils.random(16)
        key_id = f"sign_{base64.urlsafe_b64encode(key_id_bytes).decode()}"
        return (key_id, private_key_b64, public_key_b64)