Add comprehensive test suite and fix error handling in MCP server
Tests (70 tests, all passing): - TestEncode: 12 tests — text, file, encryption, channels, strategies, bit depth, compression, default paths, error cases - TestDecode: 12 tests — auto-detect, manual config, encryption, wrong password, binary payloads, strategy-specific behavior - TestAnalyze: 5 tests — clean/stegged images, full mode, channel fields - TestCapacity: 4 tests — default, single channel, high bits, errors - TestDetect: 3 tests — clean, stegged, missing image - TestChunks: 7 tests — tEXt/zTXt/iTXt/private chunk inject+read - TestInjectExif: 6 tests — fields, custom JSON, validation - TestInjectionFilename: 4 tests — templates, count, channels - TestJailbreakTemplates: 2 tests — listing and preview truncation - TestAnalysisTools: 7 tests — rs_analysis, chi_square, unicode steg, unknown action, missing file - TestCryptoStatus: 2 tests — availability and method listing - TestIntegration: 6 tests — full pipeline, encrypted pipeline, chunk+steg coexistence, large payload, sequential strategy, EXIF MCP server fixes: - All tools now return JSON errors instead of raising exceptions - Custom JSON encoder handles numpy bool_/integer/floating/ndarray - Discovered upstream bug: spread and randomized strategies fail to decode (steg_core.py issue, not MCP server)
This commit is contained in:
parent
c337ab7095
commit
7ef3ec5368
108
mcp_server.py
108
mcp_server.py
|
|
@ -14,6 +14,8 @@ import tempfile
|
|||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -69,6 +71,27 @@ except Exception:
|
|||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class _NumpyEncoder(json.JSONEncoder):
|
||||
"""Handle numpy types that standard json can't serialize."""
|
||||
|
||||
def default(self, obj):
|
||||
if isinstance(obj, (np.bool_,)):
|
||||
return bool(obj)
|
||||
if isinstance(obj, (np.integer,)):
|
||||
return int(obj)
|
||||
if isinstance(obj, (np.floating,)):
|
||||
return float(obj)
|
||||
if isinstance(obj, np.ndarray):
|
||||
return obj.tolist()
|
||||
return super().default(obj)
|
||||
|
||||
|
||||
def _json_dumps(obj: Any, **kwargs) -> str:
|
||||
"""JSON serialize with numpy type support."""
|
||||
kwargs.setdefault("cls", _NumpyEncoder)
|
||||
return json.dumps(obj, **kwargs)
|
||||
|
||||
|
||||
def _load_image(image_path: str) -> Image.Image:
|
||||
"""Load an image from disk, raising a clear error on failure."""
|
||||
p = Path(image_path).expanduser().resolve()
|
||||
|
|
@ -137,18 +160,21 @@ def stegg_encode(
|
|||
Returns:
|
||||
JSON with output_path, payload_bytes, capacity info, and encryption status.
|
||||
"""
|
||||
image = _load_image(image_path)
|
||||
try:
|
||||
image = _load_image(image_path)
|
||||
except (FileNotFoundError, OSError) as e:
|
||||
return _json_dumps({"error": str(e)})
|
||||
|
||||
# Resolve payload
|
||||
if payload_file:
|
||||
p = Path(payload_file).expanduser().resolve()
|
||||
if not p.exists():
|
||||
return json.dumps({"error": f"Payload file not found: {p}"})
|
||||
return _json_dumps({"error": f"Payload file not found: {p}"})
|
||||
payload = p.read_bytes()
|
||||
elif payload_text:
|
||||
payload = payload_text.encode("utf-8")
|
||||
else:
|
||||
return json.dumps({"error": "Provide payload_text or payload_file"})
|
||||
return _json_dumps({"error": "Provide payload_text or payload_file"})
|
||||
|
||||
out = _resolve_output(output_path or "", image_path)
|
||||
|
||||
|
|
@ -162,7 +188,7 @@ def stegg_encode(
|
|||
|
||||
capacity = calculate_capacity(image, config)
|
||||
if len(payload) > capacity["usable_bytes"]:
|
||||
return json.dumps({
|
||||
return _json_dumps({
|
||||
"error": f"Payload too large: {len(payload)} bytes > {capacity['usable_bytes']} available",
|
||||
"capacity": capacity["human"],
|
||||
})
|
||||
|
|
@ -174,7 +200,7 @@ def stegg_encode(
|
|||
|
||||
encode(image, payload, config, str(out))
|
||||
|
||||
return json.dumps({
|
||||
return _json_dumps({
|
||||
"output_path": str(out),
|
||||
"payload_bytes": len(payload),
|
||||
"capacity": capacity["human"],
|
||||
|
|
@ -218,7 +244,10 @@ def stegg_decode(
|
|||
JSON with extracted text (UTF-8) or hex preview for binary data,
|
||||
plus byte count, config detected, and output_path if saved.
|
||||
"""
|
||||
image = _load_image(image_path)
|
||||
try:
|
||||
image = _load_image(image_path)
|
||||
except (FileNotFoundError, OSError) as e:
|
||||
return _json_dumps({"error": str(e)})
|
||||
|
||||
config = None
|
||||
detected = False
|
||||
|
|
@ -242,10 +271,16 @@ def stegg_decode(
|
|||
seed=seed if seed else None,
|
||||
)
|
||||
|
||||
data = decode(image, config)
|
||||
try:
|
||||
data = decode(image, config)
|
||||
except (ValueError, Exception) as e:
|
||||
return _json_dumps({"error": f"Decode failed: {e}"})
|
||||
|
||||
if password and decrypt:
|
||||
data = decrypt(data, password)
|
||||
try:
|
||||
data = decrypt(data, password)
|
||||
except Exception as e:
|
||||
return _json_dumps({"error": f"Decryption failed: {e}"})
|
||||
|
||||
result: dict[str, Any] = {
|
||||
"bytes": len(data),
|
||||
|
|
@ -266,7 +301,7 @@ def stegg_decode(
|
|||
result["hex_preview"] = data[:512].hex()
|
||||
result["encoding"] = "binary"
|
||||
|
||||
return json.dumps(result)
|
||||
return _json_dumps(result)
|
||||
|
||||
|
||||
# ---- Analyze --------------------------------------------------------------
|
||||
|
|
@ -290,7 +325,10 @@ def stegg_analyze(
|
|||
JSON with channel stats, anomaly indicators, capacity estimates,
|
||||
and a verdict (normal / possible / high probability).
|
||||
"""
|
||||
image = _load_image(image_path)
|
||||
try:
|
||||
image = _load_image(image_path)
|
||||
except (FileNotFoundError, OSError) as e:
|
||||
return _json_dumps({"error": str(e)})
|
||||
analysis = analyze_image(image)
|
||||
|
||||
# Compact summary
|
||||
|
|
@ -337,7 +375,7 @@ def stegg_analyze(
|
|||
except Exception as e:
|
||||
result["full_analysis_error"] = str(e)
|
||||
|
||||
return json.dumps(result)
|
||||
return _json_dumps(result)
|
||||
|
||||
|
||||
# ---- Capacity -------------------------------------------------------------
|
||||
|
|
@ -358,10 +396,13 @@ def stegg_capacity(
|
|||
Returns:
|
||||
JSON with capacity in bytes and human-readable form.
|
||||
"""
|
||||
image = _load_image(image_path)
|
||||
try:
|
||||
image = _load_image(image_path)
|
||||
except (FileNotFoundError, OSError) as e:
|
||||
return _json_dumps({"error": str(e)})
|
||||
config = create_config(channels=channels, bits=bits_per_channel)
|
||||
cap = calculate_capacity(image, config)
|
||||
return json.dumps({
|
||||
return _json_dumps({
|
||||
"usable_bytes": cap["usable_bytes"],
|
||||
"human": cap["human"],
|
||||
"total_pixels": image.width * image.height,
|
||||
|
|
@ -387,11 +428,14 @@ def stegg_detect(
|
|||
Returns:
|
||||
JSON with detected (bool) and config details if found.
|
||||
"""
|
||||
image = _load_image(image_path)
|
||||
try:
|
||||
image = _load_image(image_path)
|
||||
except (FileNotFoundError, OSError) as e:
|
||||
return _json_dumps({"error": str(e)})
|
||||
detection = detect_encoding(image)
|
||||
if detection:
|
||||
return json.dumps({"detected": True, "config": detection})
|
||||
return json.dumps({"detected": False})
|
||||
return _json_dumps({"detected": True, "config": detection})
|
||||
return _json_dumps({"detected": False})
|
||||
|
||||
|
||||
# ---- PNG Chunk Injection --------------------------------------------------
|
||||
|
|
@ -423,7 +467,7 @@ def stegg_inject_chunk(
|
|||
"""
|
||||
p = Path(image_path).expanduser().resolve()
|
||||
if not p.exists():
|
||||
return json.dumps({"error": f"Image not found: {p}"})
|
||||
return _json_dumps({"error": f"Image not found: {p}"})
|
||||
raw = p.read_bytes()
|
||||
|
||||
if chunk_type == "iTXt":
|
||||
|
|
@ -436,7 +480,7 @@ def stegg_inject_chunk(
|
|||
out = Path(output_path).expanduser().resolve()
|
||||
out.write_bytes(modified)
|
||||
|
||||
return json.dumps({
|
||||
return _json_dumps({
|
||||
"output_path": str(out),
|
||||
"chunk_type": chunk_type,
|
||||
"keyword": keyword,
|
||||
|
|
@ -464,7 +508,7 @@ def stegg_read_chunks(
|
|||
"""
|
||||
p = Path(image_path).expanduser().resolve()
|
||||
if not p.exists():
|
||||
return json.dumps({"error": f"Image not found: {p}"})
|
||||
return _json_dumps({"error": f"Image not found: {p}"})
|
||||
raw = p.read_bytes()
|
||||
|
||||
chunks = read_png_chunks(raw)
|
||||
|
|
@ -479,7 +523,7 @@ def stegg_read_chunks(
|
|||
"offset": c.get("offset", 0),
|
||||
})
|
||||
|
||||
return json.dumps({
|
||||
return _json_dumps({
|
||||
"chunks": chunk_summary,
|
||||
"text_content": text_chunks,
|
||||
"total_chunks": len(chunks),
|
||||
|
|
@ -515,7 +559,7 @@ def stegg_inject_exif(
|
|||
"""
|
||||
p = Path(image_path).expanduser().resolve()
|
||||
if not p.exists():
|
||||
return json.dumps({"error": f"Image not found: {p}"})
|
||||
return _json_dumps({"error": f"Image not found: {p}"})
|
||||
|
||||
metadata: dict[str, str] = {}
|
||||
if comment:
|
||||
|
|
@ -531,10 +575,10 @@ def stegg_inject_exif(
|
|||
extra = json.loads(custom_fields)
|
||||
metadata.update(extra)
|
||||
except json.JSONDecodeError:
|
||||
return json.dumps({"error": "custom_fields must be valid JSON"})
|
||||
return _json_dumps({"error": "custom_fields must be valid JSON"})
|
||||
|
||||
if not metadata:
|
||||
return json.dumps({"error": "Provide at least one metadata field"})
|
||||
return _json_dumps({"error": "Provide at least one metadata field"})
|
||||
|
||||
image = Image.open(p)
|
||||
_, png_bytes = inject_metadata_pil(image, metadata)
|
||||
|
|
@ -542,7 +586,7 @@ def stegg_inject_exif(
|
|||
out = Path(output_path).expanduser().resolve()
|
||||
out.write_bytes(png_bytes)
|
||||
|
||||
return json.dumps({
|
||||
return _json_dumps({
|
||||
"output_path": str(out),
|
||||
"injected_fields": list(metadata.keys()),
|
||||
"field_count": len(metadata),
|
||||
|
|
@ -576,7 +620,7 @@ def stegg_injection_filename(
|
|||
generate_injection_filename(template, channels)
|
||||
for _ in range(count)
|
||||
]
|
||||
return json.dumps({
|
||||
return _json_dumps({
|
||||
"template": template,
|
||||
"channels": channels,
|
||||
"filenames": filenames,
|
||||
|
|
@ -599,7 +643,7 @@ def stegg_jailbreak_templates() -> str:
|
|||
for name in get_jailbreak_names():
|
||||
content = get_jailbreak_template(name)
|
||||
templates[name] = content[:120] + ("..." if len(content) > 120 else "")
|
||||
return json.dumps({"templates": templates, "count": len(templates)})
|
||||
return _json_dumps({"templates": templates, "count": len(templates)})
|
||||
|
||||
|
||||
# ---- Analysis Tools -------------------------------------------------------
|
||||
|
|
@ -624,14 +668,14 @@ def stegg_analysis_tool(
|
|||
"""
|
||||
p = Path(file_path).expanduser().resolve()
|
||||
if not p.exists():
|
||||
return json.dumps({"error": f"File not found: {p}"})
|
||||
return _json_dumps({"error": f"File not found: {p}"})
|
||||
data = p.read_bytes()
|
||||
|
||||
result = execute_action(action, data)
|
||||
|
||||
if hasattr(result, "to_dict"):
|
||||
return json.dumps(result.to_dict(), default=str)
|
||||
return json.dumps({"result": str(result)})
|
||||
return _json_dumps(result.to_dict(), default=str)
|
||||
return _json_dumps({"result": str(result)})
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
|
|
@ -645,7 +689,7 @@ def stegg_list_analysis_tools() -> str:
|
|||
JSON with sorted list of action names.
|
||||
"""
|
||||
tools = list_available_tools()
|
||||
return json.dumps({"tools": tools, "count": len(tools)})
|
||||
return _json_dumps({"tools": tools, "count": len(tools)})
|
||||
|
||||
|
||||
# ---- Crypto Status --------------------------------------------------------
|
||||
|
|
@ -662,8 +706,8 @@ def stegg_crypto_status() -> str:
|
|||
"""
|
||||
status = crypto_status()
|
||||
if isinstance(status, dict):
|
||||
return json.dumps(status)
|
||||
return json.dumps({"status": str(status)})
|
||||
return _json_dumps(status)
|
||||
return _json_dumps({"status": str(status)})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Binary file not shown.
|
After Width: | Height: | Size: 406 KiB |
|
|
@ -0,0 +1,670 @@
|
|||
"""
|
||||
Tests for the ST3GG MCP server.
|
||||
|
||||
Covers all 13 tools with positive, negative, and edge-case scenarios
|
||||
using a real carrier image (basi_team_six.png, 1024x1024 PNG).
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
# Ensure repo root is on sys.path
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
from mcp_server import (
|
||||
stegg_encode,
|
||||
stegg_decode,
|
||||
stegg_analyze,
|
||||
stegg_capacity,
|
||||
stegg_detect,
|
||||
stegg_inject_chunk,
|
||||
stegg_read_chunks,
|
||||
stegg_inject_exif,
|
||||
stegg_injection_filename,
|
||||
stegg_jailbreak_templates,
|
||||
stegg_analysis_tool,
|
||||
stegg_list_analysis_tools,
|
||||
stegg_crypto_status,
|
||||
)
|
||||
|
||||
FIXTURES = Path(__file__).parent / "fixtures"
|
||||
CARRIER = str(FIXTURES / "basi_team_six.png")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tmp_dir():
|
||||
with tempfile.TemporaryDirectory(prefix="stegg_test_") as d:
|
||||
yield Path(d)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# stegg_encode
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestEncode:
|
||||
def test_encode_text(self, tmp_dir):
|
||||
out = str(tmp_dir / "encoded.png")
|
||||
result = json.loads(stegg_encode(CARRIER, payload_text="hello world", output_path=out))
|
||||
assert "error" not in result
|
||||
assert result["payload_bytes"] > 0
|
||||
assert result["channels"] == "RGB"
|
||||
assert Path(out).exists()
|
||||
assert Path(out).stat().st_size > 0
|
||||
|
||||
def test_encode_file_payload(self, tmp_dir):
|
||||
payload_file = tmp_dir / "secret.txt"
|
||||
payload_file.write_text("file-based secret payload")
|
||||
out = str(tmp_dir / "encoded.png")
|
||||
result = json.loads(stegg_encode(CARRIER, payload_file=str(payload_file), output_path=out))
|
||||
assert "error" not in result
|
||||
assert result["payload_bytes"] > 0
|
||||
|
||||
def test_encode_with_encryption(self, tmp_dir):
|
||||
out = str(tmp_dir / "encrypted.png")
|
||||
result = json.loads(stegg_encode(CARRIER, payload_text="encrypted msg", output_path=out, password="s3cret"))
|
||||
assert "error" not in result
|
||||
assert result["encrypted"] is True
|
||||
|
||||
def test_encode_various_channels(self, tmp_dir):
|
||||
for ch in ["R", "G", "B", "RG", "RGBA"]:
|
||||
out = str(tmp_dir / f"encoded_{ch}.png")
|
||||
result = json.loads(stegg_encode(CARRIER, payload_text="test", output_path=out, channels=ch))
|
||||
assert "error" not in result, f"Failed for channel {ch}: {result}"
|
||||
assert result["channels"] == ch
|
||||
|
||||
def test_encode_various_strategies(self, tmp_dir):
|
||||
for strat in ["sequential", "interleaved", "spread", "randomized"]:
|
||||
out = str(tmp_dir / f"encoded_{strat}.png")
|
||||
seed = 42 if strat == "randomized" else 0
|
||||
result = json.loads(stegg_encode(
|
||||
CARRIER, payload_text="test", output_path=out,
|
||||
strategy=strat, seed=seed,
|
||||
))
|
||||
assert "error" not in result, f"Failed for strategy {strat}: {result}"
|
||||
|
||||
def test_encode_high_bit_depth(self, tmp_dir):
|
||||
out = str(tmp_dir / "encoded_4bit.png")
|
||||
result = json.loads(stegg_encode(CARRIER, payload_text="deep bits", output_path=out, bits_per_channel=4))
|
||||
assert "error" not in result
|
||||
assert result["bits_per_channel"] == 4
|
||||
|
||||
def test_encode_no_compression(self, tmp_dir):
|
||||
out = str(tmp_dir / "uncompressed.png")
|
||||
result = json.loads(stegg_encode(CARRIER, payload_text="no compress", output_path=out, compress=False))
|
||||
assert "error" not in result
|
||||
assert result["compressed"] is False
|
||||
|
||||
def test_encode_default_output_path(self, tmp_dir):
|
||||
# Use a copy in tmp_dir so default output lands there
|
||||
import shutil
|
||||
local_carrier = str(tmp_dir / "carrier.png")
|
||||
shutil.copy(CARRIER, local_carrier)
|
||||
result = json.loads(stegg_encode(local_carrier, payload_text="default path"))
|
||||
assert "error" not in result
|
||||
assert Path(result["output_path"]).exists()
|
||||
|
||||
def test_encode_missing_image(self):
|
||||
result = json.loads(stegg_encode("/nonexistent/image.png", payload_text="fail"))
|
||||
assert "error" in result
|
||||
|
||||
def test_encode_no_payload(self, tmp_dir):
|
||||
out = str(tmp_dir / "nopayload.png")
|
||||
result = json.loads(stegg_encode(CARRIER, output_path=out))
|
||||
assert "error" in result
|
||||
|
||||
def test_encode_payload_too_large(self, tmp_dir):
|
||||
out = str(tmp_dir / "toolarge.png")
|
||||
# 1024x1024 RGB 1-bit = ~384KB usable; send 1MB
|
||||
huge = "X" * (1024 * 1024)
|
||||
result = json.loads(stegg_encode(CARRIER, payload_text=huge, output_path=out))
|
||||
assert "error" in result
|
||||
assert "too large" in result["error"].lower() or "Payload" in result["error"]
|
||||
|
||||
def test_encode_missing_payload_file(self, tmp_dir):
|
||||
out = str(tmp_dir / "missing.png")
|
||||
result = json.loads(stegg_encode(CARRIER, payload_file="/nonexistent/file.bin", output_path=out))
|
||||
assert "error" in result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# stegg_decode
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestDecode:
|
||||
def _encode_helper(self, tmp_dir, text="roundtrip test", **kwargs):
|
||||
out = str(tmp_dir / "encoded.png")
|
||||
enc = json.loads(stegg_encode(CARRIER, payload_text=text, output_path=out, **kwargs))
|
||||
assert "error" not in enc
|
||||
return out
|
||||
|
||||
def test_decode_auto_detect(self, tmp_dir):
|
||||
encoded = self._encode_helper(tmp_dir, text="auto detect me")
|
||||
result = json.loads(stegg_decode(encoded))
|
||||
assert result["auto_detected"] is True
|
||||
assert result["text"] == "auto detect me"
|
||||
|
||||
def test_decode_manual_config(self, tmp_dir):
|
||||
encoded = self._encode_helper(tmp_dir, text="manual config", channels="RG", bits_per_channel=2)
|
||||
result = json.loads(stegg_decode(
|
||||
encoded, auto_detect=True,
|
||||
))
|
||||
assert result["text"] == "manual config"
|
||||
|
||||
def test_decode_with_password(self, tmp_dir):
|
||||
encoded = self._encode_helper(tmp_dir, text="secret msg", password="p4ss")
|
||||
result = json.loads(stegg_decode(encoded, password="p4ss"))
|
||||
assert result["text"] == "secret msg"
|
||||
|
||||
def test_decode_wrong_password(self, tmp_dir):
|
||||
encoded = self._encode_helper(tmp_dir, text="encrypted", password="correct")
|
||||
# Wrong password should error or return garbled data
|
||||
result = json.loads(stegg_decode(encoded, password="wrong"))
|
||||
if "error" in result:
|
||||
assert "failed" in result["error"].lower() or "invalid" in result["error"].lower()
|
||||
else:
|
||||
# If it decoded without error, text should NOT match
|
||||
assert result.get("text") != "encrypted"
|
||||
|
||||
def test_decode_save_to_file(self, tmp_dir):
|
||||
encoded = self._encode_helper(tmp_dir, text="save to disk")
|
||||
out_file = str(tmp_dir / "extracted.bin")
|
||||
result = json.loads(stegg_decode(encoded, output_path=out_file))
|
||||
assert Path(out_file).exists()
|
||||
assert Path(out_file).read_text() == "save to disk"
|
||||
|
||||
def test_decode_binary_payload(self, tmp_dir):
|
||||
# Encode a binary file
|
||||
bin_file = tmp_dir / "binary.bin"
|
||||
bin_file.write_bytes(bytes(range(256)))
|
||||
out = str(tmp_dir / "binary_encoded.png")
|
||||
enc = json.loads(stegg_encode(CARRIER, payload_file=str(bin_file), output_path=out))
|
||||
assert "error" not in enc
|
||||
|
||||
extracted = str(tmp_dir / "extracted.bin")
|
||||
result = json.loads(stegg_decode(out, output_path=extracted))
|
||||
assert Path(extracted).read_bytes() == bytes(range(256))
|
||||
|
||||
def test_decode_missing_image(self):
|
||||
result = json.loads(stegg_decode("/nonexistent/image.png"))
|
||||
assert "error" in result
|
||||
|
||||
def test_decode_interleaved_auto_detect(self, tmp_dir):
|
||||
"""Interleaved strategy supports auto-detection of the STEG header."""
|
||||
text = "interleaved-auto"
|
||||
encoded = str(tmp_dir / "enc_interleaved.png")
|
||||
enc = json.loads(stegg_encode(
|
||||
CARRIER, payload_text=text, output_path=encoded,
|
||||
strategy="interleaved",
|
||||
))
|
||||
assert "error" not in enc
|
||||
dec = json.loads(stegg_decode(encoded))
|
||||
assert dec["text"] == text
|
||||
|
||||
def test_decode_sequential_strategy_manual_config(self, tmp_dir):
|
||||
"""Sequential strategy requires manual config (auto-detect only scans interleaved)."""
|
||||
text = "strategy-sequential"
|
||||
encoded = str(tmp_dir / "enc_sequential.png")
|
||||
enc = json.loads(stegg_encode(
|
||||
CARRIER, payload_text=text, output_path=encoded,
|
||||
strategy="sequential",
|
||||
))
|
||||
assert "error" not in enc
|
||||
dec = json.loads(stegg_decode(
|
||||
encoded, auto_detect=False, strategy="sequential",
|
||||
))
|
||||
assert dec["text"] == text
|
||||
|
||||
def test_decode_spread_strategy_returns_error(self, tmp_dir):
|
||||
"""Spread and randomized strategies have known upstream decode bugs (steg_core).
|
||||
|
||||
The MCP server should return a JSON error, not crash.
|
||||
"""
|
||||
encoded = str(tmp_dir / "enc_spread.png")
|
||||
enc = json.loads(stegg_encode(
|
||||
CARRIER, payload_text="spread test", output_path=encoded,
|
||||
strategy="spread",
|
||||
))
|
||||
assert "error" not in enc # Encoding should succeed
|
||||
dec = json.loads(stegg_decode(encoded, auto_detect=False, strategy="spread"))
|
||||
# Upstream bug: spread decode fails, but server should return error gracefully
|
||||
assert "error" in dec
|
||||
|
||||
def test_decode_randomized_strategy_returns_error(self, tmp_dir):
|
||||
"""Randomized strategy has known upstream decode issues.
|
||||
|
||||
The MCP server should return a JSON error, not crash.
|
||||
"""
|
||||
encoded = str(tmp_dir / "enc_randomized.png")
|
||||
enc = json.loads(stegg_encode(
|
||||
CARRIER, payload_text="randomized test", output_path=encoded,
|
||||
strategy="randomized", seed=42,
|
||||
))
|
||||
assert "error" not in enc # Encoding should succeed
|
||||
dec = json.loads(stegg_decode(
|
||||
encoded, auto_detect=False, strategy="randomized", seed=42,
|
||||
))
|
||||
# Upstream bug: decode fails, but server should return error gracefully
|
||||
assert "error" in dec
|
||||
|
||||
def test_decode_all_channel_presets(self, tmp_dir):
|
||||
for ch in ["R", "G", "B", "RG", "RGB", "RGBA"]:
|
||||
text = f"ch-{ch}"
|
||||
encoded = str(tmp_dir / f"enc_{ch}.png")
|
||||
enc = json.loads(stegg_encode(
|
||||
CARRIER, payload_text=text, output_path=encoded, channels=ch,
|
||||
))
|
||||
assert "error" not in enc, f"Encode failed for {ch}"
|
||||
dec = json.loads(stegg_decode(encoded))
|
||||
assert dec["text"] == text, f"Decode roundtrip failed for {ch}: got {dec.get('text')}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# stegg_analyze
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestAnalyze:
|
||||
def test_analyze_clean_image(self):
|
||||
result = json.loads(stegg_analyze(CARRIER))
|
||||
assert "verdict" in result
|
||||
assert result["dimensions"]["width"] == 1024
|
||||
assert result["dimensions"]["height"] == 1024
|
||||
assert "channels" in result
|
||||
assert "capacity" in result
|
||||
|
||||
def test_analyze_stegged_image(self, tmp_dir):
|
||||
out = str(tmp_dir / "stegged.png")
|
||||
stegg_encode(CARRIER, payload_text="A" * 10000, output_path=out)
|
||||
result = json.loads(stegg_analyze(out))
|
||||
# Large payload should trigger anomaly
|
||||
assert "HIGH" in result["verdict"] or "Possible" in result["verdict"]
|
||||
|
||||
def test_analyze_full_mode(self, tmp_dir):
|
||||
out = str(tmp_dir / "stegged.png")
|
||||
stegg_encode(CARRIER, payload_text="full analysis test", output_path=out)
|
||||
result = json.loads(stegg_analyze(out, full=True))
|
||||
assert "full_analysis" in result or "full_analysis_error" in result
|
||||
|
||||
def test_analyze_missing_image(self):
|
||||
result = json.loads(stegg_analyze("/nonexistent.png"))
|
||||
assert "error" in result
|
||||
|
||||
def test_analyze_channel_fields(self):
|
||||
result = json.loads(stegg_analyze(CARRIER))
|
||||
for ch_name in ["R", "G", "B"]:
|
||||
ch = result["channels"][ch_name]
|
||||
assert "mean" in ch
|
||||
assert "std" in ch
|
||||
assert "lsb_zeros_pct" in ch
|
||||
assert "lsb_ones_pct" in ch
|
||||
assert "chi_square_indicator" in ch
|
||||
assert "anomaly" in ch
|
||||
assert ch["anomaly"] in ("normal", "slight", "HIGH")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# stegg_capacity
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestCapacity:
|
||||
def test_capacity_default(self):
|
||||
result = json.loads(stegg_capacity(CARRIER))
|
||||
assert result["usable_bytes"] > 0
|
||||
assert result["total_pixels"] == 1024 * 1024
|
||||
assert "human" in result
|
||||
|
||||
def test_capacity_single_channel(self):
|
||||
result = json.loads(stegg_capacity(CARRIER, channels="R", bits_per_channel=1))
|
||||
result_rgb = json.loads(stegg_capacity(CARRIER, channels="RGB", bits_per_channel=1))
|
||||
# Single channel should be roughly 1/3 of RGB
|
||||
assert result["usable_bytes"] < result_rgb["usable_bytes"]
|
||||
|
||||
def test_capacity_high_bits(self):
|
||||
result_1 = json.loads(stegg_capacity(CARRIER, bits_per_channel=1))
|
||||
result_4 = json.loads(stegg_capacity(CARRIER, bits_per_channel=4))
|
||||
assert result_4["usable_bytes"] > result_1["usable_bytes"]
|
||||
|
||||
def test_capacity_missing_image(self):
|
||||
result = json.loads(stegg_capacity("/nonexistent.png"))
|
||||
assert "error" in result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# stegg_detect
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestDetect:
|
||||
def test_detect_clean_image(self):
|
||||
result = json.loads(stegg_detect(CARRIER))
|
||||
assert result["detected"] is False
|
||||
|
||||
def test_detect_stegged_image(self, tmp_dir):
|
||||
out = str(tmp_dir / "stegged.png")
|
||||
stegg_encode(CARRIER, payload_text="detectable", output_path=out)
|
||||
result = json.loads(stegg_detect(out))
|
||||
assert result["detected"] is True
|
||||
assert "config" in result
|
||||
|
||||
def test_detect_missing_image(self):
|
||||
result = json.loads(stegg_detect("/nonexistent.png"))
|
||||
assert "error" in result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# stegg_inject_chunk / stegg_read_chunks
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestChunks:
|
||||
def test_inject_and_read_text_chunk(self, tmp_dir):
|
||||
out = str(tmp_dir / "chunked.png")
|
||||
inj = json.loads(stegg_inject_chunk(
|
||||
CARRIER, out, chunk_type="tEXt", keyword="Comment", text="injected comment",
|
||||
))
|
||||
assert "error" not in inj
|
||||
assert Path(out).exists()
|
||||
|
||||
read = json.loads(stegg_read_chunks(out))
|
||||
assert read["text_content"]["Comment"] == "injected comment"
|
||||
|
||||
def test_inject_compressed_chunk(self, tmp_dir):
|
||||
out = str(tmp_dir / "compressed_chunk.png")
|
||||
inj = json.loads(stegg_inject_chunk(
|
||||
CARRIER, out, chunk_type="zTXt", keyword="Description",
|
||||
text="compressed description data", compressed=True,
|
||||
))
|
||||
assert "error" not in inj
|
||||
|
||||
def test_inject_itxt_chunk(self, tmp_dir):
|
||||
out = str(tmp_dir / "itxt.png")
|
||||
inj = json.loads(stegg_inject_chunk(
|
||||
CARRIER, out, chunk_type="iTXt", keyword="Author", text="ST3GG",
|
||||
))
|
||||
assert "error" not in inj
|
||||
|
||||
def test_inject_private_chunk(self, tmp_dir):
|
||||
out = str(tmp_dir / "private.png")
|
||||
inj = json.loads(stegg_inject_chunk(
|
||||
CARRIER, out, chunk_type="stEg", keyword="", text="private data",
|
||||
))
|
||||
assert "error" not in inj
|
||||
|
||||
def test_read_chunks_structure(self):
|
||||
result = json.loads(stegg_read_chunks(CARRIER))
|
||||
assert "chunks" in result
|
||||
assert "total_chunks" in result
|
||||
assert result["total_chunks"] > 0
|
||||
# Should have at minimum IHDR, IDAT, IEND
|
||||
types = [c["type"] for c in result["chunks"]]
|
||||
assert "IHDR" in types
|
||||
assert "IEND" in types
|
||||
|
||||
def test_read_chunks_missing_image(self):
|
||||
result = json.loads(stegg_read_chunks("/nonexistent.png"))
|
||||
assert "error" in result
|
||||
|
||||
def test_inject_chunk_missing_image(self, tmp_dir):
|
||||
out = str(tmp_dir / "fail.png")
|
||||
result = json.loads(stegg_inject_chunk("/nonexistent.png", out, text="fail"))
|
||||
assert "error" in result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# stegg_inject_exif
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestInjectExif:
|
||||
def test_inject_comment_and_author(self, tmp_dir):
|
||||
out = str(tmp_dir / "exif.png")
|
||||
result = json.loads(stegg_inject_exif(
|
||||
CARRIER, out, comment="PoC comment", author="red-team",
|
||||
))
|
||||
assert "error" not in result
|
||||
assert result["field_count"] == 2
|
||||
assert Path(out).exists()
|
||||
|
||||
# Verify by reading chunks
|
||||
chunks = json.loads(stegg_read_chunks(out))
|
||||
assert chunks["text_content"].get("Comment") == "PoC comment"
|
||||
assert chunks["text_content"].get("Author") == "red-team"
|
||||
|
||||
def test_inject_custom_fields(self, tmp_dir):
|
||||
out = str(tmp_dir / "custom_exif.png")
|
||||
custom = json.dumps({"Software": "ST3GG-MCP", "X-Custom": "payload"})
|
||||
result = json.loads(stegg_inject_exif(CARRIER, out, custom_fields=custom))
|
||||
assert "error" not in result
|
||||
assert result["field_count"] == 2
|
||||
|
||||
def test_inject_all_fields(self, tmp_dir):
|
||||
out = str(tmp_dir / "all_fields.png")
|
||||
result = json.loads(stegg_inject_exif(
|
||||
CARRIER, out,
|
||||
comment="c", author="a", description="d", title="t",
|
||||
))
|
||||
assert result["field_count"] == 4
|
||||
|
||||
def test_inject_no_fields(self, tmp_dir):
|
||||
out = str(tmp_dir / "empty.png")
|
||||
result = json.loads(stegg_inject_exif(CARRIER, out))
|
||||
assert "error" in result
|
||||
|
||||
def test_inject_invalid_custom_json(self, tmp_dir):
|
||||
out = str(tmp_dir / "bad_json.png")
|
||||
result = json.loads(stegg_inject_exif(CARRIER, out, custom_fields="not json"))
|
||||
assert "error" in result
|
||||
|
||||
def test_inject_exif_missing_image(self, tmp_dir):
|
||||
out = str(tmp_dir / "fail.png")
|
||||
result = json.loads(stegg_inject_exif("/nonexistent.png", out, comment="fail"))
|
||||
assert "error" in result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# stegg_injection_filename
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestInjectionFilename:
|
||||
def test_single_filename(self):
|
||||
result = json.loads(stegg_injection_filename())
|
||||
assert len(result["filenames"]) == 1
|
||||
assert result["template"] == "universal_decoder"
|
||||
|
||||
def test_multiple_filenames(self):
|
||||
result = json.loads(stegg_injection_filename(count=5))
|
||||
assert len(result["filenames"]) == 5
|
||||
# Should all be unique (randomized)
|
||||
assert len(set(result["filenames"])) == 5
|
||||
|
||||
def test_various_templates(self):
|
||||
templates = ["chatgpt_decoder", "claude_decoder", "gemini_decoder",
|
||||
"system_override", "universal_decoder"]
|
||||
for t in templates:
|
||||
result = json.loads(stegg_injection_filename(template=t))
|
||||
assert "error" not in result
|
||||
assert result["template"] == t
|
||||
assert len(result["filenames"]) == 1
|
||||
|
||||
def test_custom_channels(self):
|
||||
result = json.loads(stegg_injection_filename(channels="RGBA"))
|
||||
assert result["channels"] == "RGBA"
|
||||
assert "RGBA" in result["filenames"][0]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# stegg_jailbreak_templates
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestJailbreakTemplates:
|
||||
def test_list_templates(self):
|
||||
result = json.loads(stegg_jailbreak_templates())
|
||||
assert result["count"] > 0
|
||||
assert "templates" in result
|
||||
assert isinstance(result["templates"], dict)
|
||||
|
||||
def test_template_previews_truncated(self):
|
||||
result = json.loads(stegg_jailbreak_templates())
|
||||
for name, preview in result["templates"].items():
|
||||
# Previews should be <= 123 chars (120 + "...")
|
||||
assert len(preview) <= 123, f"Template {name} preview too long: {len(preview)}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# stegg_analysis_tool / stegg_list_analysis_tools
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestAnalysisTools:
|
||||
def test_list_tools(self):
|
||||
result = json.loads(stegg_list_analysis_tools())
|
||||
assert result["count"] > 0
|
||||
assert "tools" in result
|
||||
assert isinstance(result["tools"], list)
|
||||
|
||||
def test_rs_analysis(self, tmp_dir):
|
||||
out = str(tmp_dir / "stegged.png")
|
||||
stegg_encode(CARRIER, payload_text="A" * 5000, output_path=out)
|
||||
result = json.loads(stegg_analysis_tool(out, "rs_analysis"))
|
||||
assert result["success"] is True
|
||||
|
||||
def test_png_chi_square_analysis(self):
|
||||
result = json.loads(stegg_analysis_tool(CARRIER, "png_chi_square_analysis"))
|
||||
assert result["success"] is True
|
||||
|
||||
def test_png_parse_chunks_tool(self):
|
||||
result = json.loads(stegg_analysis_tool(CARRIER, "png_parse_chunks"))
|
||||
assert result["success"] is True
|
||||
|
||||
def test_detect_unicode_steg(self):
|
||||
# Create a text file with zero-width chars
|
||||
import tempfile
|
||||
with tempfile.NamedTemporaryFile(suffix=".txt", delete=False, mode="w") as f:
|
||||
f.write("normal text\u200b\u200c\u200d hidden")
|
||||
f.flush()
|
||||
result = json.loads(stegg_analysis_tool(f.name, "detect_unicode_steg"))
|
||||
assert result["success"] is True
|
||||
os.unlink(f.name)
|
||||
|
||||
def test_unknown_action(self):
|
||||
result = json.loads(stegg_analysis_tool(CARRIER, "nonexistent_tool_xyz"))
|
||||
assert result["success"] is False
|
||||
assert "error" in result
|
||||
|
||||
def test_analysis_missing_file(self):
|
||||
result = json.loads(stegg_analysis_tool("/nonexistent.png", "rs_analysis"))
|
||||
assert "error" in result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# stegg_crypto_status
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestCryptoStatus:
|
||||
def test_crypto_status(self):
|
||||
result = json.loads(stegg_crypto_status())
|
||||
assert "available_methods" in result or "cryptography_available" in result
|
||||
|
||||
def test_has_methods(self):
|
||||
result = json.loads(stegg_crypto_status())
|
||||
methods = result.get("available_methods", [])
|
||||
assert len(methods) > 0
|
||||
assert "xor" in methods # Always available
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Integration: full pipeline tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestIntegration:
|
||||
"""End-to-end pipeline tests combining multiple tools."""
|
||||
|
||||
def test_full_pipeline_encode_analyze_detect_decode(self, tmp_dir):
|
||||
"""Encode -> analyze (should detect anomaly) -> detect (should find header) -> decode."""
|
||||
msg = "Full pipeline integration test with Basi Team Six!"
|
||||
out = str(tmp_dir / "pipeline.png")
|
||||
|
||||
# Encode
|
||||
enc = json.loads(stegg_encode(CARRIER, payload_text=msg, output_path=out))
|
||||
assert "error" not in enc
|
||||
|
||||
# Analyze — should show anomaly
|
||||
ana = json.loads(stegg_analyze(out))
|
||||
assert "verdict" in ana
|
||||
|
||||
# Detect — should find header
|
||||
det = json.loads(stegg_detect(out))
|
||||
assert det["detected"] is True
|
||||
|
||||
# Decode — should recover exact message
|
||||
dec = json.loads(stegg_decode(out))
|
||||
assert dec["text"] == msg
|
||||
|
||||
def test_encrypted_pipeline(self, tmp_dir):
|
||||
"""Encode encrypted -> detect -> decode with correct password."""
|
||||
msg = "encrypted pipeline test"
|
||||
pw = "Bas1T3amS1x!"
|
||||
out = str(tmp_dir / "encrypted_pipeline.png")
|
||||
|
||||
enc = json.loads(stegg_encode(CARRIER, payload_text=msg, output_path=out, password=pw))
|
||||
assert enc["encrypted"] is True
|
||||
|
||||
dec = json.loads(stegg_decode(out, password=pw))
|
||||
assert dec["text"] == msg
|
||||
|
||||
def test_chunk_injection_preserves_steg_data(self, tmp_dir):
|
||||
"""Encode data, then inject chunk — original steg data should survive."""
|
||||
msg = "steg survives chunk injection"
|
||||
stegged = str(tmp_dir / "stegged.png")
|
||||
chunked = str(tmp_dir / "chunked.png")
|
||||
|
||||
stegg_encode(CARRIER, payload_text=msg, output_path=stegged)
|
||||
stegg_inject_chunk(stegged, chunked, keyword="Comment", text="metadata")
|
||||
|
||||
# Chunk should be readable
|
||||
chunks = json.loads(stegg_read_chunks(chunked))
|
||||
assert chunks["text_content"]["Comment"] == "metadata"
|
||||
|
||||
# Steg data should still be decodable
|
||||
dec = json.loads(stegg_decode(chunked))
|
||||
assert dec["text"] == msg
|
||||
|
||||
def test_large_payload_near_capacity(self, tmp_dir):
|
||||
"""Encode near-capacity payload and verify roundtrip."""
|
||||
cap = json.loads(stegg_capacity(CARRIER, channels="RGBA", bits_per_channel=2))
|
||||
# Use 80% of capacity to stay safe after compression overhead
|
||||
size = int(cap["usable_bytes"] * 0.6)
|
||||
payload = "X" * size
|
||||
out = str(tmp_dir / "large.png")
|
||||
|
||||
enc = json.loads(stegg_encode(
|
||||
CARRIER, payload_text=payload, output_path=out,
|
||||
channels="RGBA", bits_per_channel=2,
|
||||
))
|
||||
assert "error" not in enc
|
||||
|
||||
dec = json.loads(stegg_decode(out))
|
||||
assert dec["text"] == payload
|
||||
|
||||
def test_sequential_strategy_roundtrip(self, tmp_dir):
|
||||
"""Sequential strategy encode/decode roundtrip with manual config."""
|
||||
msg = "sequential pipeline test"
|
||||
out = str(tmp_dir / "seq.png")
|
||||
|
||||
enc = json.loads(stegg_encode(CARRIER, payload_text=msg, output_path=out, strategy="sequential"))
|
||||
assert "error" not in enc
|
||||
|
||||
dec = json.loads(stegg_decode(out, auto_detect=False, strategy="sequential"))
|
||||
assert dec["text"] == msg
|
||||
|
||||
def test_exif_injection_pipeline(self, tmp_dir):
|
||||
"""Inject EXIF, then read chunks to verify, then analyze."""
|
||||
out = str(tmp_dir / "exif_pipeline.png")
|
||||
stegg_inject_exif(CARRIER, out, comment="pipeline test", author="0xmoose")
|
||||
|
||||
chunks = json.loads(stegg_read_chunks(out))
|
||||
assert chunks["text_content"]["Comment"] == "pipeline test"
|
||||
assert chunks["text_content"]["Author"] == "0xmoose"
|
||||
|
||||
ana = json.loads(stegg_analyze(out))
|
||||
assert "verdict" in ana
|
||||
Loading…
Reference in New Issue