95 lines
2.9 KiB
Python
95 lines
2.9 KiB
Python
import secrets
|
|
from pydantic import BaseModel
|
|
from src.utils import (
|
|
generate_random_list, generate_random_matrix, xor_lists,
|
|
)
|
|
|
|
|
|
class UserEncipherKeys(BaseModel):
|
|
alpha_key: list[list[int]]
|
|
set_key: list[int]
|
|
pass_key: list[int]
|
|
mask_key: list[int]
|
|
# TODO: add these
|
|
# height
|
|
# width
|
|
# max_nkode_len
|
|
|
|
@staticmethod
|
|
def new_user_encipher_keys(height: int, width: int, set_values: list[int]):
|
|
assert len(set_values) == width
|
|
|
|
set_key = generate_random_list(width)
|
|
set_key = xor_lists(set_key, set_values)
|
|
|
|
return UserEncipherKeys(
|
|
alpha_key=generate_random_matrix(width, height),
|
|
pass_key=generate_random_list(height),
|
|
mask_key=generate_random_list(height),
|
|
set_key=set_key,
|
|
)
|
|
|
|
|
|
class CustomerAttribute(BaseModel):
|
|
attr_val: int
|
|
set_val: int
|
|
|
|
|
|
class CustomerInterface(BaseModel):
|
|
# base_interface: { set0: [(attr0, set0) ... (attrN, set0)], ... setM: [(attr0, setM) ... (attrN, setM)]}
|
|
base_interface: dict[int, CustomerAttribute]
|
|
height: int
|
|
width: int
|
|
|
|
def __init__(self, **data):
|
|
super().__init__(**data)
|
|
self._set_index_lookup: dict[int, int] = {}
|
|
self._attr_index_lookup: dict[CustomerAttribute: (int, int)] = {}
|
|
self._map_base_interface()
|
|
|
|
def _map_base_interface(self):
|
|
for set_idx, set_val in enumerate(self.base_interface.keys()):
|
|
attrs = self.base_interface[set_val]
|
|
self._set_index_lookup[set_val] = set_idx
|
|
for attr_idx, attr in enumerate(attrs):
|
|
self._attr_index_lookup[attr] = attr_idx
|
|
|
|
@staticmethod
|
|
def new_interface(height: int, width: int):
|
|
assert (height <= 256)
|
|
assert (width <= 256)
|
|
|
|
base_interface = {}
|
|
possible_set_vals = list(range(256))
|
|
for w in range(width):
|
|
cur_set = secrets.choice(possible_set_vals)
|
|
possible_attr_values = list(range(256))
|
|
base_interface[cur_set] = []
|
|
for h in range(height):
|
|
cur_attr = secrets.choice(possible_attr_values)
|
|
base_interface[cur_set].append(CustomerAttribute(attr_val=cur_attr, set_val=cur_set))
|
|
possible_attr_values.remove(cur_attr)
|
|
possible_set_vals.remove(cur_set)
|
|
|
|
return CustomerInterface(
|
|
base_interface=base_interface,
|
|
height=height,
|
|
width=width,
|
|
)
|
|
|
|
def get_set_values(self) -> list[int]:
|
|
return list(self.base_interface.keys())
|
|
|
|
def get_attr_index(self, attr: CustomerAttribute) -> int:
|
|
assert(attr in self._attr_index_lookup.keys())
|
|
return self._attr_index_lookup[attr]
|
|
|
|
def get_set_index(self, set_val: int) -> int:
|
|
assert(set_val in self.get_set_values())
|
|
return self._set_index_lookup[set_val]
|
|
|
|
|
|
class EncipheredNKode(BaseModel):
|
|
code: str
|
|
mask: str
|