From 03d7c6bafc770af2594d1975d22ddb9e63e25867 Mon Sep 17 00:00:00 2001 From: Arpit Jain Date: Wed, 30 Sep 2026 02:38:07 -0400 Subject: [PATCH] Make the payload encoder's -key argument usable Signed-off-by: Arpit Jain --- app/utility/payload_encoder.py | 6 +++-- tests/utility/test_payload_encoder.py | 39 +++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) create mode 100644 tests/utility/test_payload_encoder.py diff --git a/app/utility/payload_encoder.py b/app/utility/payload_encoder.py index 9e8f2a8a1..422bd5e25 100644 --- a/app/utility/payload_encoder.py +++ b/app/utility/payload_encoder.py @@ -10,7 +10,7 @@ Additionally, payload_encoder.py can be used from the command-line to add a new encoded payload. ``` -python /path/to/payload_encoder.py input_file output_file +python /path/to/payload_encoder.py [-key KEY] input_file output_file ``` NOTE: In order for the server to detect the availability of an encoded payload, the payload file's @@ -44,7 +44,9 @@ def xor_file(input_file, output_file=None, key=None): if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('-key', default=DEFAULT_KEY) + # A str indexes to str, which cannot be XOR'ed with an int. The server path in + # file_svc encodes its key the same way before calling xor_bytes. + parser.add_argument('-key', default=DEFAULT_KEY, type=lambda key: key.encode()) parser.add_argument('input') parser.add_argument('output') diff --git a/tests/utility/test_payload_encoder.py b/tests/utility/test_payload_encoder.py new file mode 100644 index 000000000..a33d0dc89 --- /dev/null +++ b/tests/utility/test_payload_encoder.py @@ -0,0 +1,39 @@ +import subprocess +import sys + +from pathlib import Path + +from app.utility.payload_encoder import DEFAULT_KEY, xor_bytes + + +ENCODER = Path(__file__).parents[2] / 'app' / 'utility' / 'payload_encoder.py' + + +def _run_encoder(args): + return subprocess.run([sys.executable, str(ENCODER)] + args, capture_output=True, text=True) + + +def test_xor_bytes_accepts_a_byte_string_key(): + assert xor_bytes(b'\x00\x01', key=b'\x0f\x0f') == b'\x0f\x0e' + + +def test_command_line_uses_the_default_key(tmp_path): + payload = tmp_path / 'payload' + payload.write_bytes(b'hello payload') + encoded = tmp_path / 'payload.xored' + + result = _run_encoder([str(payload), str(encoded)]) + + assert result.returncode == 0, result.stderr + assert encoded.read_bytes() == xor_bytes(b'hello payload', key=DEFAULT_KEY) + + +def test_command_line_key_argument_is_usable(tmp_path): + payload = tmp_path / 'payload' + payload.write_bytes(b'hello payload') + encoded = tmp_path / 'payload.xored' + + result = _run_encoder(['-key', 'somekey', str(payload), str(encoded)]) + + assert result.returncode == 0, result.stderr + assert encoded.read_bytes() == xor_bytes(b'hello payload', key=b'somekey')