Files
bitwarden-desktop-agent/crypto.py
2026-03-20 02:50:59 +09:00

105 lines
3.1 KiB
Python

import base64
import hashlib
import hmac
import os
from cryptography.hazmat.primitives import padding as sym_padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from secmem import SecureBuffer, wipe
class SymmetricKey:
def __init__(self, raw: bytes | bytearray):
if len(raw) != 64:
raise ValueError(f"expected 64 bytes, got {len(raw)}")
self._secure = SecureBuffer(raw)
if isinstance(raw, bytearray):
wipe(raw)
@property
def raw(self) -> bytearray:
return self._secure.raw
@property
def enc_key(self) -> bytearray:
return self._secure.raw[:32]
@property
def mac_key(self) -> bytearray:
return self._secure.raw[32:]
def to_b64(self) -> str:
return base64.b64encode(bytes(self._secure)).decode()
def close(self):
self._secure.close()
@classmethod
def from_b64(cls, b64: str) -> "SymmetricKey":
return cls(base64.b64decode(b64))
@classmethod
def generate(cls) -> "SymmetricKey":
return cls(os.urandom(64))
def _decrypt_raw(enc_str: str, key: SymmetricKey) -> bytearray:
t, rest = enc_str.split(".", 1)
parts = rest.split("|")
iv = base64.b64decode(parts[0])
ct = base64.b64decode(parts[1])
if len(parts) > 2:
mac_got = base64.b64decode(parts[2])
expected = hmac.new(bytes(key.mac_key), iv + ct, hashlib.sha256).digest()
if not hmac.compare_digest(mac_got, expected):
raise ValueError("MAC mismatch")
dec = Cipher(algorithms.AES(bytes(key.enc_key)), modes.CBC(iv)).decryptor()
padded = dec.update(ct) + dec.finalize()
unpadder = sym_padding.PKCS7(128).unpadder()
return bytearray(unpadder.update(padded) + unpadder.finalize())
def enc_string_encrypt(plaintext: str, key: SymmetricKey) -> str:
iv = os.urandom(16)
padder = sym_padding.PKCS7(128).padder()
padded = padder.update(plaintext.encode()) + padder.finalize()
ct = Cipher(algorithms.AES(bytes(key.enc_key)), modes.CBC(iv)).encryptor()
encrypted = ct.update(padded) + ct.finalize()
mac = hmac.new(bytes(key.mac_key), iv + encrypted, hashlib.sha256).digest()
iv_b64 = base64.b64encode(iv).decode()
ct_b64 = base64.b64encode(encrypted).decode()
mac_b64 = base64.b64encode(mac).decode()
return f"2.{iv_b64}|{ct_b64}|{mac_b64}"
def enc_string_decrypt(enc_str: str, key: SymmetricKey) -> str:
raw = _decrypt_raw(enc_str, key)
result = raw.decode()
wipe(raw)
return result
def enc_string_decrypt_bytes(enc_str: str, key: SymmetricKey) -> bytearray:
return _decrypt_raw(enc_str, key)
def enc_string_to_dict(enc_str: str) -> dict:
t, rest = enc_str.split(".", 1)
parts = rest.split("|")
d = {"encryptionType": int(t), "encryptedString": enc_str}
if len(parts) >= 1:
d["iv"] = parts[0]
if len(parts) >= 2:
d["data"] = parts[1]
if len(parts) >= 3:
d["mac"] = parts[2]
return d
def dict_to_enc_string(d: dict) -> str:
if s := d.get("encryptedString"):
return s
t = d.get("encryptionType", 2)
return f"{t}.{d.get('iv', '')}|{d.get('data', '')}|{d.get('mac', '')}"