diff --git a/src/models.py b/src/models.py index 7192274..6c9b9de 100644 --- a/src/models.py +++ b/src/models.py @@ -116,6 +116,14 @@ class OuterKey(DarcKey): 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 @@ -128,6 +136,14 @@ class InnerKey(DarcKey): 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 @@ -138,7 +154,14 @@ class SubstitutionKey(DarcKey): matrix = [cls._random_mask(width, max_value) for _ in range(height)] return SubstitutionKey( matrix=matrix, - max_value=max_value + ) + + @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, ) diff --git a/test/test_darc.py b/test/test_darc.py index a283fc7..420e0f2 100644 --- a/test/test_darc.py +++ b/test/test_darc.py @@ -12,7 +12,7 @@ import pytest ) def test_darc(height, width, message_len): alphabet = SubstitutionKey.init_matrix(width, height, 255) - medium = SubstitutionKey(matrix=[[0 for _ in range(width)] for _ in range(height)]) + medium = SubstitutionKey.init_identity_matrix(width, height) server_keys, mutual_keys, client_keys = generate_keys(height, width) original_message = OuterKey.init_matrix(height).matrix[0]