refactor use numpy in user_cipher.py
This commit is contained in:
@@ -9,7 +9,7 @@ from src.utils import xor_lists
|
||||
class Customer:
|
||||
customer_id: UUID
|
||||
nkode_policy: NKodePolicy
|
||||
customer_cipher: CustomerCipher
|
||||
cipher: CustomerCipher
|
||||
users: dict[str, User]
|
||||
|
||||
# TODO: validate policy and keypad size don't conflict
|
||||
@@ -23,16 +23,16 @@ class Customer:
|
||||
if username not in self.users:
|
||||
raise ValueError(f"User '{username}' does not exist")
|
||||
|
||||
numb_of_keys = self.customer_cipher.keypad_size.numb_of_keys
|
||||
numb_of_keys = self.cipher.keypad_size.numb_of_keys
|
||||
if not all(0 <= key_idx < numb_of_keys for key_idx in selected_keys):
|
||||
raise ValueError(f"Invalid key indices. Must be between 0 and {numb_of_keys - 1}")
|
||||
|
||||
passcode_len = len(selected_keys)
|
||||
user = self.users[username]
|
||||
|
||||
passcode_set_vals = user.user_keys.decipher_mask(
|
||||
user.enciphered_passcode.mask, self.customer_cipher.set_key, passcode_len)
|
||||
set_vals_idx = [self.customer_cipher.get_set_index(set_val) for set_val in passcode_set_vals]
|
||||
passcode_set_vals = user.cipher.decipher_mask(
|
||||
user.enciphered_passcode.mask, self.cipher.set_key, passcode_len)
|
||||
set_vals_idx = [self.cipher.get_set_index(set_val) for set_val in passcode_set_vals]
|
||||
|
||||
presumed_selected_attributes_idx = []
|
||||
for idx in range(passcode_len):
|
||||
@@ -41,20 +41,20 @@ class Customer:
|
||||
selected_attr_idx = user.user_keypad.get_attr_idx_by_keynumb_setidx(key_numb, set_idx)
|
||||
presumed_selected_attributes_idx.append(selected_attr_idx)
|
||||
|
||||
enciphered_attr = user.user_keys.encipher_salt_hash_code(presumed_selected_attributes_idx, self.customer_cipher)
|
||||
enciphered_attr = user.cipher.encipher_salt_hash_code(presumed_selected_attributes_idx, self.cipher)
|
||||
if enciphered_attr != user.enciphered_passcode.code:
|
||||
return False
|
||||
|
||||
if user.renew:
|
||||
user.refresh_passcode(presumed_selected_attributes_idx, self.customer_cipher)
|
||||
user.refresh_passcode(presumed_selected_attributes_idx, self.cipher)
|
||||
return True
|
||||
|
||||
def renew_keys(self) -> bool:
|
||||
old_attrs = self.customer_cipher.prop_key.copy()
|
||||
old_sets = self.customer_cipher.set_key.copy()
|
||||
self.customer_cipher.renew()
|
||||
new_attrs = self.customer_cipher.prop_key
|
||||
new_sets = self.customer_cipher.set_key
|
||||
old_attrs = self.cipher.prop_key.copy()
|
||||
old_sets = self.cipher.set_key.copy()
|
||||
self.cipher.renew()
|
||||
new_attrs = self.cipher.prop_key
|
||||
new_sets = self.cipher.set_key
|
||||
|
||||
attrs_xor = xor_lists(new_attrs, old_attrs)
|
||||
set_xor = xor_lists(new_sets, old_sets)
|
||||
@@ -66,7 +66,7 @@ class Customer:
|
||||
def valid_new_nkode(self, passcode_attr_idx: list[int]) -> bool:
|
||||
nkode_len = len(passcode_attr_idx)
|
||||
passcode_set_values = [
|
||||
self.customer_cipher.get_prop_set_val(self.customer_cipher.prop_key[attr_idx]) for attr_idx in passcode_attr_idx
|
||||
self.cipher.get_prop_set_val(self.cipher.prop_key[attr_idx]) for attr_idx in passcode_attr_idx
|
||||
]
|
||||
distinct_sets = len(set(passcode_set_values))
|
||||
distinct_attributes = len(set(passcode_attr_idx))
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
from dataclasses import dataclass, field
|
||||
from uuid import UUID, uuid4
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
from src.customer import Customer
|
||||
from src.models import NKodePolicy, KeypadSize
|
||||
from src.user import User
|
||||
@@ -19,7 +18,7 @@ class NKodeAPI:
|
||||
def create_new_customer(self, keypad_size: KeypadSize, nkode_policy: NKodePolicy) -> UUID:
|
||||
new_customer = Customer(
|
||||
customer_id=uuid4(),
|
||||
customer_cipher=CustomerCipher.create(keypad_size),
|
||||
cipher=CustomerCipher.create(keypad_size),
|
||||
users={},
|
||||
nkode_policy=nkode_policy
|
||||
)
|
||||
@@ -30,7 +29,7 @@ class NKodeAPI:
|
||||
if customer_id not in self.customers.keys():
|
||||
raise ValueError(f"Customer with ID '{customer_id}' does not exist")
|
||||
customer = self.customers[customer_id]
|
||||
login_keypad = UserKeypad.create(customer.customer_cipher.keypad_size)
|
||||
login_keypad = UserKeypad.create(customer.cipher.keypad_size)
|
||||
set_keypad = login_keypad.sign_up_keypad()
|
||||
new_session = UserSignupSession(
|
||||
session_id=uuid4(),
|
||||
@@ -76,15 +75,15 @@ class NKodeAPI:
|
||||
customer = self.customers[customer_id]
|
||||
passcode = self.signup_sessions[session_id].deduce_passcode(confirm_key_entry)
|
||||
new_user_keys = UserCipher.create(
|
||||
customer.customer_cipher.keypad_size,
|
||||
customer.customer_cipher.set_key,
|
||||
customer.cipher.keypad_size,
|
||||
customer.cipher.set_key,
|
||||
customer.nkode_policy.max_nkode_len
|
||||
)
|
||||
enciphered_passcode = new_user_keys.encipher_nkode(passcode, customer.customer_cipher)
|
||||
enciphered_passcode = new_user_keys.encipher_nkode(passcode, customer.cipher)
|
||||
new_user = User(
|
||||
username=username,
|
||||
enciphered_passcode=enciphered_passcode,
|
||||
user_keys=new_user_keys,
|
||||
cipher=new_user_keys,
|
||||
user_keypad=self.signup_sessions[session_id].login_keypad,
|
||||
)
|
||||
self.customers[customer_id].add_new_user(new_user)
|
||||
|
||||
14
src/user.py
14
src/user.py
@@ -10,20 +10,20 @@ from src.utils import xor_lists
|
||||
class User:
|
||||
username: str
|
||||
enciphered_passcode: EncipheredNKode
|
||||
user_keys: UserCipher
|
||||
cipher: UserCipher
|
||||
user_keypad: UserKeypad
|
||||
renew: bool = field(default=False)
|
||||
|
||||
def renew_keys(self, sets_xor: list[int], attrs_xor: list[int]):
|
||||
def renew_keys(self, set_xor: list[int], prop_xor: list[int]):
|
||||
self.renew = True
|
||||
self.user_keys.set_key = xor_lists(self.user_keys.set_key, sets_xor)
|
||||
self.user_keys.prop_key = xor_lists(self.user_keys.prop_key, attrs_xor)
|
||||
self.cipher.set_key = xor_lists(self.cipher.set_key, set_xor)
|
||||
self.cipher.prop_key = xor_lists(self.cipher.prop_key, prop_xor)
|
||||
|
||||
def refresh_passcode(self, passcode_attr_idx: list[int], customer_attributes: CustomerCipher):
|
||||
self.user_keys = UserCipher.create(
|
||||
self.cipher = UserCipher.create(
|
||||
customer_attributes.keypad_size,
|
||||
customer_attributes.set_key,
|
||||
self.user_keys.max_nkode_len
|
||||
self.cipher.max_nkode_len
|
||||
)
|
||||
self.enciphered_passcode = self.user_keys.encipher_nkode(passcode_attr_idx, customer_attributes)
|
||||
self.enciphered_passcode = self.cipher.encipher_nkode(passcode_attr_idx, customer_attributes)
|
||||
self.renew = False
|
||||
|
||||
@@ -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])
|
||||
Reference in New Issue
Block a user