Files
pydarc/src/models.py
2024-07-04 09:13:48 -05:00

502 lines
15 KiB
Python

import secrets
from enum import Enum
from pydantic import BaseModel
class DarcKeyType(str, Enum):
outer_shuffle = "outer_shuffle"
inner_shuffle = "inner_shuffle"
substitution = "substitution"
other = "other"
class DarcKey(BaseModel):
matrix: list[list[int]]
key_type: DarcKeyType
@classmethod
def _random_shuffle(cls, array: list[int]) -> list[int]:
for i in range(len(array) - 1, 0, -1):
j = cls._random_number_in_range(0, i)
array[i], array[j] = array[j], array[i]
return array
@classmethod
def _random_mask(cls, array_len: int, max_val: int) -> list[int]:
return [cls._random_number_in_range(0, max_val) for _ in range(array_len)]
@staticmethod
def _random_number_in_range(min_val: int, max_val: int) -> int:
""" originally secure_rand """
return min_val + secrets.randbelow(max_val - min_val + 1)
def __lshift__(self, other):
assert len(self.matrix) == len(other.matrix)
assert len(self.matrix[0]) == len(other.matrix[0])
assert other.key_type == DarcKeyType.inner_shuffle or other.key_type == DarcKeyType.outer_shuffle
output = []
for i, el in enumerate(self.matrix):
row = []
for j, el2 in enumerate(el):
other_el = other.matrix[i][j]
row.append(self.matrix[i][other_el])
output.append(row)
return DarcKey(
matrix=output,
key_type=self.key_type
)
def __lt__(self, other):
assert len(other.matrix) == 1
assert len(self.matrix) == len(other.matrix[0])
assert other.key_type == DarcKeyType.outer_shuffle
output = []
for i, _ in enumerate(self.matrix):
output.append(self.matrix[other.matrix[0][i]])
return DarcKey(
matrix=output,
key_type=self.key_type
)
def __xor__(self, other):
assert len(self.matrix) == len(other.matrix)
assert len(self.matrix[0]) == len(other.matrix[0])
assert self.key_type == DarcKeyType.substitution
assert other.key_type == DarcKeyType.substitution
output = []
for i, el in enumerate(self.matrix):
row = []
for j, el2 in enumerate(el):
row.append(el2 ^ other.matrix[i][j])
output.append(row)
return DarcKey(
matrix=output,
key_type=self.key_type
)
def __invert__(self):
assert self.key_type == DarcKeyType.inner_shuffle or self.key_type == DarcKeyType.outer_shuffle
output = []
for i, el in enumerate(self.matrix):
row = [(j, el2) for j, el2 in enumerate(el)]
row.sort(key=lambda x: x[1])
output.append([el2[0] for el2 in row])
return DarcKey(
matrix=output,
key_type=self.key_type
)
def __eq__(self, other) -> bool:
if self.key_type != other.key_type:
return False
if len(self.matrix) != len(other.matrix):
return False
if len(self.matrix[0]) != len(other.matrix[0]):
return False
for i in range(len(self.matrix)):
for j in range(len(self.matrix[0])):
if self.matrix[i][j] != other.matrix[i][j]:
return False
return True
class OuterKey(DarcKey):
key_type: DarcKeyType = DarcKeyType.outer_shuffle
@classmethod
def init_matrix(cls, height: int):
""" originally buildRandom """
matrix = [cls._random_shuffle(list(range(height)))]
return OuterKey(
matrix=matrix,
)
@classmethod
def init_identity_matrix(cls, height: int):
""" originally buildRandom """
matrix = [list(range(height))]
return OuterKey(
matrix=matrix,
)
class InnerKey(DarcKey):
key_type: DarcKeyType = DarcKeyType.inner_shuffle
@classmethod
def init_matrix(cls, width: int, height: int):
""" originally buildRandom """
matrix = [cls._random_shuffle(list(range(width))) for _ in range(height)]
return InnerKey(
matrix=matrix,
)
@classmethod
def init_identity_matrix(cls, width: int, height: int):
""" originally buildRandom """
matrix = [list(range(width)) for _ in range(height)]
return InnerKey(
matrix=matrix,
)
class SubstitutionKey(DarcKey):
key_type: DarcKeyType = DarcKeyType.substitution
@classmethod
def init_matrix(cls, width: int, height: int, max_value: int = 255):
""" originally buildScramble """
matrix = [cls._random_mask(width, max_value) for _ in range(height)]
return SubstitutionKey(
matrix=matrix,
)
@classmethod
def init_identity_matrix(cls, width: int, height: int):
""" originally buildScramble """
matrix = [[0 for _ in range(width)] for _ in range(height)]
return SubstitutionKey(
matrix=matrix,
)
class Mask(DarcKey):
key_type: DarcKeyType = DarcKeyType.substitution
@classmethod
def init_matrix(cls, width: int, height: int, col_mask: int):
""" originally buildScramble """
max_value = 255
mask = cls._random_mask(width, max_value)
matrix = [[0 for _ in range(width)] for _ in range(height)]
if col_mask < height:
matrix[col_mask] = mask
return Mask(
matrix=matrix,
)
class ClientEphemeralDataKeys(BaseModel):
outer_key_1: OuterKey
outer_key_2: OuterKey
alpha_key: SubstitutionKey
inner_key_1: InnerKey
@classmethod
def random_init(cls, height: int, width: int):
return ClientEphemeralDataKeys(
outer_key_1=OuterKey.init_matrix(height),
outer_key_2=OuterKey.init_matrix(height),
alpha_key=SubstitutionKey.init_matrix(width, height, 255),
inner_key_1=InnerKey.init_matrix(width, height),
)
class ClientEphemeralMediumKeys(BaseModel):
outer_key_1: OuterKey
outer_key_2: OuterKey
inner_key_1: InnerKey
alpha_key: SubstitutionKey
@classmethod
def random_init(cls, height: int, width: int):
return ClientEphemeralMediumKeys(
outer_key_1=OuterKey.init_matrix(height),
outer_key_2=OuterKey.init_matrix(height),
inner_key_1=InnerKey.init_matrix(width, height),
alpha_key=SubstitutionKey.init_matrix(width, height, 255)
)
class ClientPersistentDataKeys(BaseModel):
alpha_key: SubstitutionKey
inner_key_1: InnerKey
outer_key_1: OuterKey
outer_key_2: OuterKey
@classmethod
def random_init(cls, height: int, width: int):
return ClientPersistentDataKeys(
alpha_key=SubstitutionKey.init_matrix(width, height, 255),
inner_key_1=InnerKey.init_matrix(width, height),
outer_key_1=OuterKey.init_matrix(height),
outer_key_2=OuterKey.init_matrix(height),
)
class ClientPersistentMediumKeys(BaseModel):
outer_key_1: OuterKey
outer_key_2: OuterKey
inner_key_1: InnerKey
alpha_key: SubstitutionKey
@classmethod
def random_init(cls, height: int, width: int):
return ClientPersistentMediumKeys(
outer_key_1=OuterKey.init_matrix(height),
outer_key_2=OuterKey.init_matrix(height),
inner_key_1=InnerKey.init_matrix(width, height),
alpha_key=SubstitutionKey.init_matrix(width, height, 255)
)
class ServerEphemeralDataKeys(BaseModel):
outer_key_1: OuterKey
alpha_key: SubstitutionKey
inner_key_1: InnerKey
@classmethod
def random_init(cls, height: int, width: int):
return ServerEphemeralDataKeys(
outer_key_1=OuterKey.init_matrix(height),
alpha_key=SubstitutionKey.init_matrix(width, height, 255),
inner_key_1=InnerKey.init_matrix(width, height),
)
class ServerEphemeralMediumKeys(BaseModel):
alpha_key: SubstitutionKey
@classmethod
def random_init(cls, height: int, width: int):
return ServerEphemeralMediumKeys(
alpha_key=SubstitutionKey.init_matrix(width, height, 255)
)
class ServerPersistentDataKeys(BaseModel):
outer_key_1: OuterKey
alpha_key: SubstitutionKey
inner_key_1: InnerKey
@classmethod
def random_init(cls, height: int, width: int):
return ServerPersistentDataKeys(
outer_key_1=OuterKey.init_matrix(height),
alpha_key=SubstitutionKey.init_matrix(width, height, 255),
inner_key_1=InnerKey.init_matrix(width, height),
)
class ServerPersistentMediumKeys(BaseModel):
alpha_key: SubstitutionKey
@classmethod
def random_init(cls, height: int, width: int):
return ServerPersistentMediumKeys(
alpha_key=SubstitutionKey.init_matrix(width, height, 255)
)
class MutualEphemeralDataKeys(BaseModel):
alpha_key: SubstitutionKey
inner_key_1: InnerKey
outer_key_1: OuterKey
outer_key_2: OuterKey
outer_key_3: OuterKey
@classmethod
def random_init(cls, height: int, width: int):
return MutualEphemeralDataKeys(
alpha_key=SubstitutionKey.init_matrix(width, height, 255),
inner_key_1=InnerKey.init_matrix(width, height),
outer_key_1=OuterKey.init_matrix(height),
outer_key_2=OuterKey.init_matrix(height),
outer_key_3=OuterKey.init_matrix(height),
)
class MutualEphemeralMediumKeys(BaseModel):
alpha_key: SubstitutionKey
outer_key_1: OuterKey
outer_key_2: OuterKey
outer_key_3: OuterKey
outer_key_4: OuterKey
outer_key_5: OuterKey
outer_key_6: OuterKey
inner_key_1: InnerKey
inner_key_4: InnerKey
inner_key_3: InnerKey
@classmethod
def random_init(cls, height: int, width: int):
return MutualEphemeralMediumKeys(
alpha_key=SubstitutionKey.init_matrix(width, height, 255),
outer_key_1=OuterKey.init_matrix(height),
outer_key_2=OuterKey.init_matrix(height),
outer_key_3=OuterKey.init_matrix(height),
outer_key_4=OuterKey.init_matrix(height),
outer_key_5=OuterKey.init_matrix(height),
outer_key_6=OuterKey.init_matrix(height),
inner_key_1=InnerKey.init_matrix(width, height),
inner_key_4=InnerKey.init_matrix(width, height),
inner_key_3=InnerKey.init_matrix(width, height),
)
class MutualPersistentDataKeys(BaseModel):
outer_key_1: OuterKey
outer_key_2: OuterKey
outer_key_3: OuterKey
alpha_key: SubstitutionKey
inner_key_1: InnerKey
inner_key_2: InnerKey
@classmethod
def random_init(cls, height: int, width: int):
return MutualPersistentDataKeys(
outer_key_1=OuterKey.init_matrix(height),
outer_key_2=OuterKey.init_matrix(height),
outer_key_3=OuterKey.init_matrix(height),
alpha_key=SubstitutionKey.init_matrix(width, height, 255),
inner_key_1=InnerKey.init_matrix(width, height),
inner_key_2=InnerKey.init_matrix(width, height),
)
class MutualPersistentMediumKeys(BaseModel):
alpha_key: SubstitutionKey
outer_key_1: OuterKey
outer_key_2: OuterKey
outer_key_3: OuterKey
outer_key_4: OuterKey
outer_key_5: OuterKey
outer_key_6: OuterKey
inner_key_1: InnerKey
inner_key_2: InnerKey
inner_key_3: InnerKey
inner_key_4: InnerKey
@classmethod
def random_init(cls, height: int, width: int):
return MutualPersistentMediumKeys(
alpha_key=SubstitutionKey.init_matrix(width, height, 255),
outer_key_1=OuterKey.init_matrix(height),
outer_key_2=OuterKey.init_matrix(height),
outer_key_3=OuterKey.init_matrix(height),
outer_key_4=OuterKey.init_matrix(height),
outer_key_5=OuterKey.init_matrix(height),
outer_key_6=OuterKey.init_matrix(height),
inner_key_1=InnerKey.init_matrix(width, height),
inner_key_2=InnerKey.init_matrix(width, height),
inner_key_3=InnerKey.init_matrix(width, height),
inner_key_4=InnerKey.init_matrix(width, height),
)
class ServerPersistentKeys(BaseModel):
data: ServerPersistentDataKeys
medium: ServerPersistentMediumKeys
@classmethod
def random_init(cls, height: int, width: int):
return ServerPersistentKeys(
data=ServerPersistentDataKeys.random_init(height, width),
medium=ServerPersistentMediumKeys.random_init(height, width)
)
class ClientPersistentKeys(BaseModel):
data: ClientPersistentDataKeys
medium: ClientPersistentMediumKeys
@classmethod
def random_init(cls, height: int, width: int):
return ClientPersistentKeys(
data=ClientPersistentDataKeys.random_init(height, width),
medium=ClientPersistentMediumKeys.random_init(height, width)
)
class MutualPersistentKeys(BaseModel):
data: MutualPersistentDataKeys
medium: MutualPersistentMediumKeys
@classmethod
def random_init(cls, height: int, width: int):
return MutualPersistentKeys(
data=MutualPersistentDataKeys.random_init(height, width),
medium=MutualPersistentMediumKeys.random_init(height, width)
)
class ServerEphemeralKeys(BaseModel):
data: ServerEphemeralDataKeys
medium: ServerEphemeralMediumKeys
@classmethod
def random_init(cls, height: int, width: int):
return ServerEphemeralKeys(
data=ServerEphemeralDataKeys.random_init(height, width),
medium=ServerEphemeralMediumKeys.random_init(height, width)
)
class MutualEphemeralKeys(BaseModel):
data: MutualEphemeralDataKeys
medium: MutualEphemeralMediumKeys
@classmethod
def random_init(cls, height: int, width: int):
return MutualEphemeralKeys(
data=MutualEphemeralDataKeys.random_init(height, width),
medium=MutualEphemeralMediumKeys.random_init(height, width)
)
class ClientEphemeralKeys(BaseModel):
data: ClientEphemeralDataKeys
medium: ClientEphemeralMediumKeys
@classmethod
def random_init(cls, height: int, width: int):
return ClientEphemeralKeys(
data=ClientEphemeralDataKeys.random_init(height, width),
medium=ClientEphemeralMediumKeys.random_init(height, width)
)
class ServerKeys(BaseModel):
persistent: ServerPersistentKeys
ephemeral: ServerEphemeralKeys
@classmethod
def random_init(cls, height: int, width: int):
return ServerKeys(
persistent=ServerPersistentKeys.random_init(height, width),
ephemeral=ServerEphemeralKeys.random_init(height, width)
)
class MutualKeys(BaseModel):
persistent: MutualPersistentKeys
ephemeral: MutualEphemeralKeys
@classmethod
def random_init(cls, height: int, width: int):
return MutualKeys(
persistent=MutualPersistentKeys.random_init(height, width),
ephemeral=MutualEphemeralKeys.random_init(height, width)
)
class ClientKeys(BaseModel):
persistent: ClientPersistentKeys
ephemeral: ClientEphemeralKeys
@classmethod
def random_init(cls, height: int, width: int):
return ClientKeys(
persistent=ClientPersistentKeys.random_init(height, width),
ephemeral=ClientEphemeralKeys.random_init(height, width)
)