245 lines
7.2 KiB
Python
245 lines
7.2 KiB
Python
import argparse
|
|
import json
|
|
import math
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from util import is_valid_zip, select_zip
|
|
|
|
"""
|
|
Checks if all submission output the correct decrypted data.
|
|
Directory structure should be as follows:
|
|
|
|
submissions
|
|
├── <team1>
|
|
| ├── mem_layout.json [optional]
|
|
| └── <savename>.zip
|
|
├── <team2>
|
|
| ├── mem_layout.json [optional]
|
|
| └── <savename>.zip
|
|
...
|
|
|
|
group
|
|
├── sBox
|
|
├── key
|
|
├── data_encrypted
|
|
└── data_decrypted
|
|
|
|
mem_layout.json must contain the keys 'sbox', 'key' and 'data', with the respective addresses as values.
|
|
If the file does not exist, the user will be prompted to enter these addresses. The data
|
|
will then be saved as 'mem_layout.json', to be used in subsequent calls.
|
|
|
|
If more than one simulator ZIP file is present, the user will also be prompted to choose.
|
|
"""
|
|
|
|
SIMULATOR_PATH = Path("./minimax_simulator-2.0.0-cli.jar")
|
|
|
|
# helper strings
|
|
OK = "\033[32mOK\033[0m"
|
|
ERROR = "\033[31mERROR\033[0m"
|
|
|
|
|
|
def ul(s: str) -> str:
|
|
"""
|
|
Adds ansi escape sequences to underline string.
|
|
"""
|
|
return f"\033[4m{s}\033[0m"
|
|
|
|
|
|
def compare_result(actual_file: Path, expected_file: Path) -> bool:
|
|
"""
|
|
Compares the file at path 'actual_file' with the file at path 'expected_file' bytewise.
|
|
If the actual file is larger than the expected file, additional bytes will be ignored.
|
|
"""
|
|
try:
|
|
with open(actual_file, "rb") as actual, open(expected_file, "rb") as expected:
|
|
actual_bytes = actual.read()
|
|
expected_bytes = expected.read()
|
|
except FileNotFoundError:
|
|
return False
|
|
|
|
return actual_bytes[: len(expected_bytes)] == expected_bytes
|
|
|
|
|
|
def create_mem_layout() -> dict:
|
|
"""
|
|
Prompts the user to enter addresses for the sbox, key and data.
|
|
Returns a dictionary with these three keys.
|
|
"""
|
|
mem_layout = dict()
|
|
for key in ("sbox", "key", "data"):
|
|
while True:
|
|
try:
|
|
value = input(f"{key} address (prefix hex numbers with '0x'): ")
|
|
base = 16 if value.startswith("0x") else 10
|
|
value = int(value, base=base)
|
|
break
|
|
except ValueError:
|
|
print("Adress is not an integer.")
|
|
except KeyboardInterrupt:
|
|
sys.exit("Aborted.")
|
|
mem_layout[key] = value
|
|
return mem_layout
|
|
|
|
|
|
def evaluate(
|
|
zip_file: Path,
|
|
sbox_file: Path,
|
|
key_file: Path,
|
|
data_file: Path,
|
|
result_file: Path,
|
|
mem_layout: dict,
|
|
simulator: Path,
|
|
) -> None:
|
|
"""
|
|
Runs the minimax simulator on the given input. The resulting file is saved in 'result_file'.
|
|
"""
|
|
args = [
|
|
"java",
|
|
"-jar",
|
|
simulator,
|
|
zip_file,
|
|
"--import-file",
|
|
sbox_file,
|
|
"--import-from",
|
|
mem_layout["sbox"],
|
|
"--import-file",
|
|
key_file,
|
|
"--import-from",
|
|
mem_layout["key"],
|
|
"--import-file",
|
|
data_file,
|
|
"--import-from",
|
|
mem_layout["data"],
|
|
"--export-file",
|
|
result_file,
|
|
"--export-from",
|
|
mem_layout.get("result", mem_layout["data"]),
|
|
"--export-to",
|
|
mem_layout.get("result", mem_layout["data"]) + math.ceil(os.path.getsize(data_file) / 4),
|
|
]
|
|
|
|
if "result" in mem_layout:
|
|
print("Decryption was not done in-place.")
|
|
|
|
args = [str(arg) for arg in args]
|
|
# subprocess.run requires all arguments to be strings
|
|
|
|
print(f"Running simulator, storing result in '{result_file}'")
|
|
print("\033[38;5;245m")
|
|
|
|
subprocess.run(args, stdout=sys.stdout, stderr=sys.stderr)
|
|
|
|
print("\033[0m")
|
|
|
|
|
|
if __name__ == "__main__": # only run if executed as script
|
|
parser = argparse.ArgumentParser(
|
|
description="Runs the simulator on all projects found in the submissions directory."
|
|
)
|
|
parser.add_argument(
|
|
"submissions",
|
|
type=Path,
|
|
help="Submissions root directory",
|
|
)
|
|
parser.add_argument(
|
|
"group",
|
|
type=Path,
|
|
help="Group directory, contains all project files",
|
|
)
|
|
parser.add_argument(
|
|
"-e",
|
|
"--file-extension",
|
|
dest="file_ext",
|
|
type=str,
|
|
help="Result file extension",
|
|
)
|
|
parser.add_argument("-j", "--jar", dest="jar", type=Path, help="Simulator jar file")
|
|
args = parser.parse_args()
|
|
|
|
# Load teams
|
|
teams = [e for e in args.submissions.iterdir() if e.is_dir()]
|
|
|
|
print(f"The following teams were found:")
|
|
for team in teams:
|
|
print(f"* {team.name}")
|
|
|
|
# Check directory structure
|
|
for filename in ("sBox", "key", "data_encrypted", "data_decrypted"):
|
|
if not (args.group / filename).exists():
|
|
sys.exit(f"Group project file '{filename}' is missing.")
|
|
|
|
# Check file extension
|
|
if args.file_ext is None:
|
|
args.file_ext = ""
|
|
elif not args.file_ext.startswith("."):
|
|
args.file_ext = f".{args.file_ext}"
|
|
|
|
# Update simulator path if given
|
|
simulator = SIMULATOR_PATH if args.jar is None else args.jar
|
|
|
|
# Store evaluation results
|
|
expected_file = args.group / "data_decrypted"
|
|
results = []
|
|
|
|
# Evaluate each team
|
|
for team in teams:
|
|
print(ul(f"Evaluating team '{team.name}'"))
|
|
|
|
# load memory layout file if available (otherwise create and store it)
|
|
try:
|
|
with open(team / "mem_layout.json", "r") as mem_layout_file:
|
|
mem_layout = json.load(mem_layout_file)
|
|
except FileNotFoundError:
|
|
mem_layout = create_mem_layout()
|
|
with open(team / "mem_layout.json", "w") as mem_layout_file:
|
|
json.dump(mem_layout, mem_layout_file)
|
|
|
|
# Select project file (if more there is more than one zip file)
|
|
zip_files = [file for file in team.glob("*.zip") if is_valid_zip(file)]
|
|
zip_file = zip_files[0] if len(zip_files) == 1 else select_zip(zip_files)
|
|
|
|
# check memory layout and convert hexadecimal addresses
|
|
for key in ("sbox", "key", "data"):
|
|
if key not in mem_layout:
|
|
sys.exit(f"memory_layout does not contain key '{key}'")
|
|
try:
|
|
addr = int(mem_layout[key])
|
|
except ValueError:
|
|
try:
|
|
addr = int(mem_layout[key], base=16)
|
|
except ValueError:
|
|
sys.exit(f"Invalid address '{mem_layout[key]}' for key '{key}'")
|
|
mem_layout[key] = addr
|
|
|
|
# evaluate team
|
|
sbox_file = args.group / "sBox"
|
|
key_file = args.group / "key"
|
|
data_file = args.group / "data_encrypted"
|
|
result_file = team / f"data_decrypted{args.file_ext}"
|
|
|
|
evaluate(
|
|
zip_file,
|
|
sbox_file,
|
|
key_file,
|
|
data_file,
|
|
result_file,
|
|
mem_layout,
|
|
simulator=simulator,
|
|
)
|
|
|
|
# Compare against expected result
|
|
results.append(compare_result(result_file, expected_file))
|
|
|
|
# Print result summary
|
|
print("Summary:")
|
|
for team, result in zip(teams, results):
|
|
if result is True:
|
|
print(f"[{OK}] - {team.name}")
|
|
else:
|
|
print(f"[{ERROR}] - {team.name}")
|
|
|
|
print()
|