greydgl--pentestgpt/tests/unit/test_flag_detection.py

"""Tests for flag detection patterns.

Unit tests for the flag detection regex patterns used in PentestAgent.
"""

import re

import pytest

# Flag patterns from pentestgpt/core/controller.py
FLAG_PATTERNS = [
    r"flag\{[^\}]+\}",  # flag{...}
    r"FLAG\{[^\}]+\}",  # FLAG{...}
    r"HTB\{[^\}]+\}",  # HTB{...}
    r"CTF\{[^\}]+\}",  # CTF{...}
    r"[A-Za-z0-9_]+\{[^\}]+\}",  # Generic CTF format
    r"\b[a-f0-9]{32}\b",  # 32-char hex (HTB user/root flags)
]


def detect_flags(text: str) -> list[str]:
    """Detect potential flags in text using regex patterns."""
    flags = []
    for pattern in FLAG_PATTERNS:
        matches = re.finditer(pattern, text, re.IGNORECASE)
        for match in matches:
            flag = match.group(0)
            if flag not in flags:
                flags.append(flag)
    return flags


@pytest.mark.unit
class TestFlagDetection:
    """Tests for flag detection patterns."""

    def test_detect_flag_lowercase(self):
        """Test detecting lowercase flag{...} pattern."""
        text = "The flag is flag{s3cr3t_p4ssw0rd}"
        flags = detect_flags(text)
        assert "flag{s3cr3t_p4ssw0rd}" in flags

    def test_detect_flag_uppercase(self):
        """Test detecting uppercase FLAG{...} pattern."""
        text = "Found: FLAG{ADMIN_ACCESS}"
        flags = detect_flags(text)
        assert "FLAG{ADMIN_ACCESS}" in flags

    def test_detect_htb_flag(self):
        """Test detecting HTB{...} pattern."""
        text = "HackTheBox flag: HTB{h4ck_th3_pl4n3t}"
        flags = detect_flags(text)
        assert "HTB{h4ck_th3_pl4n3t}" in flags

    def test_detect_ctf_flag(self):
        """Test detecting CTF{...} pattern."""
        text = "Capture the flag: CTF{y0u_w1n}"
        flags = detect_flags(text)
        assert "CTF{y0u_w1n}" in flags

    def test_detect_custom_ctf_format(self):
        """Test detecting custom CTF format like picoCTF{...}."""
        text = "Here is the flag: picoCTF{custom_flag_here}"
        flags = detect_flags(text)
        assert "picoCTF{custom_flag_here}" in flags

    def test_detect_32_char_hex(self):
        """Test detecting 32-character hex flags (HTB user/root)."""
        text = "User flag: 0123456789abcdef0123456789abcdef"
        flags = detect_flags(text)
        assert "0123456789abcdef0123456789abcdef" in flags

    def test_detect_32_char_hex_mixed_case(self):
        """Test detecting 32-char hex with mixed case."""
        text = "Root flag: aBcDeF0123456789aBcDeF0123456789"
        flags = detect_flags(text)
        # The pattern should match (case insensitive)
        assert any("abcdef0123456789abcdef0123456789" in f.lower() for f in flags)

    def test_detect_multiple_flags(self):
        """Test detecting multiple flags in one text."""
        text = """
        Found these flags:
        - flag{first_flag}
        - HTB{second_flag}
        - 0123456789abcdef0123456789abcdef
        """
        flags = detect_flags(text)
        assert len(flags) >= 3
        assert "flag{first_flag}" in flags
        assert "HTB{second_flag}" in flags
        assert "0123456789abcdef0123456789abcdef" in flags

    def test_no_false_positives_short_hex(self):
        """Test that short hex strings are not detected as flags."""
        text = "Hash: abc123 and ID: 12345678"
        flags = detect_flags(text)
        # Should not detect short hex strings
        assert "abc123" not in flags
        assert "12345678" not in flags

    def test_no_false_positives_31_char_hex(self):
        """Test that 31-char hex strings are not detected as flags."""
        text = "Almost: 0123456789abcdef0123456789abcde"  # 31 chars
        flags = detect_flags(text)
        # Should not include 31-char strings
        assert not any(len(f) == 31 and f.isalnum() for f in flags)

    def test_no_false_positives_33_char_hex(self):
        """Test that 33-char hex strings are not detected as flags."""
        text = "Too long: 0123456789abcdef0123456789abcdefa"  # 33 chars
        flags = detect_flags(text)
        # The pattern should not match the full 33-char string
        # (may match first 32 chars depending on pattern)
        assert "0123456789abcdef0123456789abcdefa" not in flags

    def test_flag_with_special_content(self):
        """Test detecting flags with special characters in content."""
        text = "flag{c0mp13x_fl4g_w1th_und3rsc0r3s}"
        flags = detect_flags(text)
        assert "flag{c0mp13x_fl4g_w1th_und3rsc0r3s}" in flags

    def test_flag_in_json_context(self):
        """Test detecting flags in JSON-like output."""
        text = '{"result": "success", "flag": "flag{json_flag_123}"}'
        flags = detect_flags(text)
        assert "flag{json_flag_123}" in flags

    def test_flag_in_command_output(self):
        """Test detecting flags in command output context."""
        text = """
        cat /root/root.txt
        HTB{r00t_4cc3ss_gr4nt3d}
        """
        flags = detect_flags(text)
        assert "HTB{r00t_4cc3ss_gr4nt3d}" in flags

    def test_empty_string(self):
        """Test that empty string returns no flags."""
        flags = detect_flags("")
        assert flags == []

    def test_no_flags(self):
        """Test text with no flags returns empty list."""
        text = "This is just regular text without any flags."
        flags = detect_flags(text)
        # May detect some false positives with generic pattern, but
        # main flag formats should not match
        assert not any(f.startswith(("flag{", "FLAG{", "HTB{", "CTF{")) for f in flags)

    def test_flag_case_insensitive(self):
        """Test that flag detection is case insensitive."""
        text = "FlaG{MiXeD_CaSe}"
        flags = detect_flags(text)
        # Should detect the flag
        assert len(flags) >= 1

    def test_duplicate_flags_not_repeated(self):
        """Test that duplicate flags are not returned multiple times."""
        text = "flag{duplicate} and again flag{duplicate}"
        flags = detect_flags(text)
        # Count occurrences of the duplicate flag
        count = flags.count("flag{duplicate}")
        assert count == 1