refactor use numpy in user_cipher.py

This commit is contained in:
2025-03-11 10:33:08 -05:00
parent dd0b496a21
commit 526e537586
7 changed files with 230 additions and 177 deletions

View File

@@ -2,17 +2,18 @@ import base64
import hashlib
from dataclasses import dataclass
import bcrypt
import numpy as np
from secrets import choice
from src.models import EncipheredNKode, KeypadSize
from src.customer_cipher import CustomerCipher
from src.utils import generate_random_nonrepeating_list, xor_lists, int_array_to_bytes
@dataclass
class UserCipher:
prop_key: list[int]
set_key: list[int]
pass_key: list[int]
mask_key: list[int]
prop_key: np.ndarray
set_key: np.ndarray
pass_key: np.ndarray
mask_key: np.ndarray
salt: bytes
max_nkode_len: int
@@ -21,41 +22,50 @@ class UserCipher:
if len(set_values) != keypad_size.props_per_key:
raise ValueError("Invalid set values")
set_key = generate_random_nonrepeating_list(keypad_size.props_per_key)
set_key = xor_lists(set_key, set_values)
set_values_array = np.array(set_values, dtype=np.uint16)
set_key = generate_random_nonrepeating_array(keypad_size.props_per_key)
set_key = np.bitwise_xor(set_key, set_values_array)
return UserCipher(
prop_key=generate_random_nonrepeating_list(keypad_size.props_per_key * keypad_size.numb_of_keys),
pass_key=generate_random_nonrepeating_list(max_nkode_len),
mask_key=generate_random_nonrepeating_list(max_nkode_len),
prop_key=generate_random_nonrepeating_array(keypad_size.props_per_key * keypad_size.numb_of_keys),
pass_key=generate_random_nonrepeating_array(max_nkode_len),
mask_key=generate_random_nonrepeating_array(max_nkode_len),
set_key=set_key,
salt=bcrypt.gensalt(),
max_nkode_len=max_nkode_len
)
def pad_user_mask(self, user_mask: list[int], set_vals: list[int]) -> list[int]:
def pad_user_mask(self, user_mask: list[int], set_vals: list[int]) -> np.ndarray:
if len(user_mask) >= self.max_nkode_len:
raise ValueError("User mask is too long")
padded_user_mask = user_mask.copy()
for _ in range(self.max_nkode_len - len(user_mask)):
padded_user_mask.append(choice(set_vals))
user_mask_array = np.array(user_mask, dtype=np.uint16)
set_vals_array = np.array(set_vals, dtype=np.uint16)
# Create padding of random choices from set_vals
padding_size = self.max_nkode_len - len(user_mask)
padding_indices = np.random.choice(len(set_vals), padding_size)
padding = np.array([set_vals[i] for i in padding_indices], dtype=np.uint16)
# Concatenate original mask with padding
padded_user_mask = np.concatenate([user_mask_array, padding])
return padded_user_mask
@staticmethod
def encode_base64_str(data: list[int]) -> str:
def encode_base64_str(data: np.ndarray) -> str:
return base64.b64encode(int_array_to_bytes(data)).decode("utf-8")
@staticmethod
def decode_base64_str(data: str) -> list[int]:
def decode_base64_str(data: str) -> np.ndarray:
byte_data = base64.b64decode(data)
int_list = []
for i in range(0, len(byte_data), 2):
int_val = int.from_bytes(byte_data[i:i + 2], byteorder='big')
int_list.append(int_val)
return int_list
return np.array(int_list, dtype=np.uint16)
def _hash_passcode(self, passcode: list[int]) -> str:
def _hash_passcode(self, passcode: np.ndarray) -> str:
passcode_bytes = int_array_to_bytes(passcode)
passcode_digest = base64.b64encode(hashlib.sha256(passcode_bytes).digest())
hashed_data = bcrypt.hashpw(passcode_digest, self.salt)
@@ -66,10 +76,10 @@ class UserCipher:
passcode_prop_idx: list[int],
customer_cipher: CustomerCipher
) -> EncipheredNKode:
passcode_attrs = [customer_cipher.prop_key[idx] for idx in passcode_prop_idx]
passcode_sets = [customer_cipher.get_prop_set_val(attr) for attr in passcode_attrs]
mask = self.encipher_mask(passcode_sets, customer_cipher)
passcode_prop_idx_array = np.array(passcode_prop_idx, dtype=np.uint16)
passcode_attrs = np.array([customer_cipher.prop_key[idx] for idx in passcode_prop_idx_array], dtype=np.uint16)
passcode_sets = np.array([customer_cipher.get_prop_set_val(attr) for attr in passcode_attrs], dtype=np.uint16)
mask = self.encipher_mask(passcode_sets.tolist(), customer_cipher)
code = self.encipher_salt_hash_code(passcode_prop_idx, customer_cipher)
return EncipheredNKode(
code=code,
@@ -81,12 +91,15 @@ class UserCipher:
passcode_prop_idx: list[int],
customer_prop: CustomerCipher,
) -> str:
passcode_len = len(passcode_prop_idx)
passcode_attrs = [customer_prop.prop_key[idx] for idx in passcode_prop_idx]
passcode_prop_idx_array = np.array(passcode_prop_idx, dtype=np.uint16)
passcode_len = len(passcode_prop_idx_array)
passcode_attrs = np.array([customer_prop.prop_key[idx] for idx in passcode_prop_idx_array], dtype=np.uint16)
passcode_cipher = self.pass_key.copy()
for idx in range(passcode_len):
attr_idx = passcode_prop_idx[idx]
passcode_cipher[idx] ^= self.prop_key[attr_idx] ^ passcode_attrs[idx]
attr_idx = passcode_prop_idx_array[idx]
passcode_cipher[idx] = passcode_cipher[idx] ^ self.prop_key[attr_idx] ^ passcode_attrs[idx]
return self._hash_passcode(passcode_cipher)
def encipher_mask(
@@ -95,19 +108,47 @@ class UserCipher:
customer_attributes: CustomerCipher
) -> str:
padded_passcode_sets = self.pad_user_mask(passcode_sets, customer_attributes.set_key)
set_idx = [customer_attributes.get_set_index(set_val) for set_val in padded_passcode_sets]
mask_set_keys = [self.set_key[idx] for idx in set_idx]
ciphered_mask = xor_lists(mask_set_keys, padded_passcode_sets)
ciphered_mask = xor_lists(ciphered_mask, self.mask_key)
# Get indices of set values
set_idx = np.array([customer_attributes.get_set_index(set_val) for set_val in padded_passcode_sets],
dtype=np.uint16)
mask_set_keys = np.array([self.set_key[idx] for idx in set_idx], dtype=np.uint16)
# XOR operations
ciphered_mask = np.bitwise_xor(mask_set_keys, padded_passcode_sets)
ciphered_mask = np.bitwise_xor(ciphered_mask, self.mask_key)
mask = self.encode_base64_str(ciphered_mask)
return mask
def decipher_mask(self, mask: str, set_vals: list, passcode_len: int) -> list[int]:
set_vals_array = np.array(set_vals, dtype=np.uint16)
decoded_mask = self.decode_base64_str(mask)
deciphered_mask = xor_lists(decoded_mask, self.mask_key)
set_key_rand_component = xor_lists(set_vals, self.set_key)
deciphered_mask = np.bitwise_xor(decoded_mask, self.mask_key)
set_key_rand_component = np.bitwise_xor(set_vals_array, self.set_key)
passcode_sets = []
for set_cipher in deciphered_mask[:passcode_len]:
set_idx = set_key_rand_component.index(set_cipher)
# Find index where values match
set_idx = np.where(set_key_rand_component == set_cipher)[0][0]
passcode_sets.append(set_vals[set_idx])
return passcode_sets
# NumPy utility functions to replace the existing ones
def generate_random_nonrepeating_array(array_len: int, min_val: int = 0, max_val: int = 2 ** 16) -> np.ndarray:
if max_val - min_val < array_len:
raise ValueError("Range of values is less than the array length requested")
# Generate array of random unique integers
return np.random.choice(
np.arange(min_val, max_val, dtype=np.uint16),
size=array_len,
replace=False
)
def int_array_to_bytes(int_arr: np.ndarray, byte_size: int = 2) -> bytes:
return b"".join([int(num).to_bytes(byte_size, byteorder='big') for num in int_arr])