diff --git a/tests/test_forensics.py b/tests/test_forensics.py index 2246125..068fd4e 100644 --- a/tests/test_forensics.py +++ b/tests/test_forensics.py @@ -354,9 +354,67 @@ def test_exif_cli_options(): result_exif = runner.invoke(cli, ["inspect", "-e", str(test_file)]) assert result_exif.exit_code == 0 assert "Metadata: mock_exif_image.png" not in result_exif.output - assert "EXIF Metadata: mock_exif_image.png" in result_exif.output + assert "EXIF Metadata" in result_exif.output + assert "mock_exif_image.png" in result_exif.output assert "Make" in result_exif.output assert "TEST" in result_exif.output +def test_jpeg_comment_extraction(): + """ + Verifies that get_metadata successfully parses a JPEG comment (COM marker). + """ + import struct + comment_text = b"picoCTF{test_comment_flag}" + comment_len = len(comment_text) + 2 + mock_jpeg = ( + b"\xff\xd8" + b"\xff\xfe" + + struct.pack(">H", comment_len) + + comment_text + + b"\xff\xd9" + ) + + test_file = TEST_ENV / "mock_comment.jpg" + with open(test_file, "wb") as f: + f.write(mock_jpeg) + + meta = get_metadata(test_file) + assert meta.comment == "picoCTF{test_comment_flag}" + +def test_exif_cli_options_warning(): + """ + Verifies CLI warns when no EXIF or comment is found and -e option is used. + """ + from ctf.main import cli + runner = CliRunner() + test_file = TEST_ENV / "clean_no_exif.png" + with open(test_file, "wb") as f: + f.write(b"\x89PNG\r\n\x1a\n") + + result = runner.invoke(cli, ["inspect", "-e", str(test_file)]) + assert result.exit_code == 0 + assert "No EXIF or comment metadata found" in result.output + +def test_exif_cli_comment_only(): + """ + Verifies CLI displays comment when -e is used and only comment is found. + """ + from ctf.main import cli + import struct + runner = CliRunner() + comment_text = b"flag{comment_only}" + comment_len = len(comment_text) + 2 + mock_jpeg = b"\xff\xd8\xff\xfe" + struct.pack(">H", comment_len) + comment_text + + test_file = TEST_ENV / "comment_only.jpg" + with open(test_file, "wb") as f: + f.write(mock_jpeg) + + result = runner.invoke(cli, ["inspect", "-e", str(test_file)]) + assert result.exit_code == 0 + assert "Comment" in result.output + assert "flag{comment_only}" in result.output + +