Files
pynkode/src/models.py
2024-07-11 08:53:25 -05:00

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