# tests/test_steg.py
import shutil
import subprocess
from pathlib import Path
from unittest.mock import patch, MagicMock
import click
from click.testing import CliRunner
import pytest

from ctf.steg import crack_steghide, SteghideCrackResult

def test_crack_steghide_missing_binary():
    """
    Verifies that crack_steghide returns failure if steghide is not installed.
    """
    with patch("shutil.which", return_value=None):
        res = crack_steghide(Path("somefile.jpg"), ["pass1", "pass2"])
        assert not res.success
        assert "not found on system" in res.error_message

def test_crack_steghide_file_not_found():
    """
    Verifies that crack_steghide returns failure if the input file does not exist.
    """
    with patch("shutil.which", return_value="steghide"):
        res = crack_steghide(Path("non_existent_file_12345.jpg"), ["pass1", "pass2"])
        assert not res.success
        assert "File not found" in res.error_message

def test_crack_steghide_success():
    """
    Verifies that crack_steghide returns success, the password, and payload when cracked.
    """
    wordlist = ["wrong1", "correct_pass", "wrong2"]
    
    # We want to mock subprocess.run to write a payload to the out_file (-xf)
    # when the correct password is tried.
    def mock_run(cmd, *args, **kwargs):
        # cmd: ['steghide', 'extract', '-sf', '...', '-p', password, '-xf', out_file, '-f']
        password = cmd[5]
        out_file = cmd[7]
        mock_res = MagicMock()
        if password == "correct_pass":
            mock_res.returncode = 0
            # Write a mock payload file
            Path(out_file).write_bytes(b"flag{steghide_cracked_payload}")
        else:
            mock_res.returncode = 1
        return mock_res

    with patch("shutil.which", return_value="steghide"), \
         patch("subprocess.run", side_effect=mock_run), \
         patch("pathlib.Path.exists", return_value=True):
        
        res = crack_steghide(Path("mock_image.jpg"), wordlist)
        assert res.success
        assert res.password == "correct_pass"
        assert res.payload == b"flag{steghide_cracked_payload}"
        assert res.error_message is None

def test_crack_steghide_failure():
    """
    Verifies that crack_steghide returns failure if the wordlist is exhausted.
    """
    wordlist = ["wrong1", "wrong2"]
    
    def mock_run(cmd, *args, **kwargs):
        mock_res = MagicMock()
        mock_res.returncode = 1
        return mock_res

    with patch("shutil.which", return_value="steghide"), \
         patch("subprocess.run", side_effect=mock_run), \
         patch("pathlib.Path.exists", return_value=True):
        
        res = crack_steghide(Path("mock_image.jpg"), wordlist)
        assert not res.success
        assert res.password is None
        assert res.payload is None
        assert "Password not found in wordlist" in res.error_message

def test_cli_crack_steghide_success(tmp_path):
    """
    Verifies that the crack-steghide CLI command prints success panel and previews payload.
    """
    from ctf.main import cli
    runner = CliRunner()
    
    # Create temp wordlist file
    wordlist_file = tmp_path / "words.txt"
    wordlist_file.write_text("pass1\npass2\n")
    
    mock_file = tmp_path / "mock.jpg"
    mock_file.write_text("dummy")
    
    mock_result = SteghideCrackResult(
        success=True,
        password="pass2",
        payload=b"flag{cli_steg_preview}"
    )
    
    with patch("ctf.cli_steg.crack_steghide", return_value=mock_result):
        res = runner.invoke(cli, ["steg", "crack-steghide", str(mock_file), "-w", str(wordlist_file)])
        assert res.exit_code == 0
        assert "Successfully cracked!" in res.output
        assert "pass2" in res.output
        assert "flag{cli_steg_preview}" in res.output

def test_cli_crack_steghide_success_output_file(tmp_path):
    """
    Verifies that the crack-steghide CLI command writes the payload to a file when specified.
    """
    from ctf.main import cli
    runner = CliRunner()
    
    wordlist_file = tmp_path / "words.txt"
    wordlist_file.write_text("pass1\n")
    
    mock_file = tmp_path / "mock.jpg"
    mock_file.write_text("dummy")
    
    out_file = tmp_path / "out.txt"
    
    mock_result = SteghideCrackResult(
        success=True,
        password="pass1",
        payload=b"written_file_content"
    )
    
    with patch("ctf.cli_steg.crack_steghide", return_value=mock_result):
        res = runner.invoke(cli, ["steg", "crack-steghide", str(mock_file), "-w", str(wordlist_file), "-o", str(out_file)])
        assert res.exit_code == 0
        assert "Successfully cracked!" in res.output
        assert out_file.exists()
        assert out_file.read_bytes() == b"written_file_content"

def test_cli_crack_steghide_failure(tmp_path):
    """
    Verifies that the crack-steghide CLI command displays correct error message on failure.
    """
    from ctf.main import cli
    runner = CliRunner()
    
    wordlist_file = tmp_path / "words.txt"
    wordlist_file.write_text("pass1\n")
    
    mock_file = tmp_path / "mock.jpg"
    mock_file.write_text("dummy")
    
    mock_result = SteghideCrackResult(
        success=False,
        error_message="Password not found in wordlist."
    )
    
    with patch("ctf.cli_steg.crack_steghide", return_value=mock_result):
        res = runner.invoke(cli, ["steg", "crack-steghide", str(mock_file), "-w", str(wordlist_file)])
        assert res.exit_code == 0
        assert "Cracking failed: Password not found in wordlist" in res.output