File size: 5,299 Bytes
0986c85
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
"""
Qwen Coder API - Security Test Script

Tests:
  1. Protected routes reject requests with no API key (401)
  2. Protected routes reject requests with a WRONG key (401) - this also
     re-confirms the substring-bypass bug is actually fixed (see test 3)
  3. A key that's a SUBSTRING of the real one is rejected (401) - this is
     the specific bug we found and fixed (the old `in` comparison would
     have let this one through)
  4. Protected routes accept requests with the CORRECT key (200)
  5. Rate limiting kicks in after RATE_LIMIT requests within 60s (429)
  6. Flags routes that are currently unprotected (/add_project, /cancel)
     so you don't forget they're still open

Usage:
    python test_security.py --url http://localhost:8080 --key YOUR_REAL_API_KEY

Or, to read the key from your .env file automatically:
    pip install python-dotenv
    python test_security.py --url http://localhost:8080
"""

import argparse
import requests

try:
    from dotenv import load_dotenv
    import os
    load_dotenv()
except ImportError:
    pass


PASS = "\033[92mPASS\033[0m"
FAIL = "\033[91mFAIL\033[0m"
WARN = "\033[93mWARN\033[0m"


def check(label: str, condition: bool, detail: str = ""):
    status = PASS if condition else FAIL
    print(f"  [{status}] {label}" + (f" — {detail}" if detail else ""))
    return condition


def test_no_key(base_url: str):
    print("\n1. Requests with NO API key should be rejected (401)")
    r = requests.post(f"{base_url}/add_project", json={"project": "test_no_key"})
    if r.status_code == 401:
        check("/add_project rejects missing key", True)
    else:
        print(f"  [{WARN}] /add_project has no auth protection yet (got {r.status_code}, not 401) — known open item")

    r = requests.post(f"{base_url}/chat", json={"project": "default", "message": "hi"})
    check("/chat rejects missing key", r.status_code == 401, f"got {r.status_code}")


def test_wrong_key(base_url: str, real_key: str):
    print("\n2. Requests with a WRONG key should be rejected (401)")
    r = requests.post(
        f"{base_url}/chat",
        json={"project": "default", "message": "hi"},
        headers={"X-API-Key": "definitely-not-the-real-key"},
    )
    check("/chat rejects wrong key", r.status_code == 401, f"got {r.status_code}")


def test_substring_bug(base_url: str, real_key: str):
    print("\n3. A SUBSTRING of the real key should be rejected (the original bug)")
    if len(real_key) < 4:
        print(f"  [{WARN}] Real key too short to meaningfully test substring bypass, skipping")
        return
    substring_key = real_key[: max(3, len(real_key) // 3)]
    r = requests.post(
        f"{base_url}/chat",
        json={"project": "default", "message": "hi"},
        headers={"X-API-Key": substring_key},
    )
    check(
        "/chat rejects a substring of the real key",
        r.status_code == 401,
        f"sent '{substring_key}...', got {r.status_code} (200 here would mean the old bug is back)",
    )


def test_correct_key(base_url: str, real_key: str):
    print("\n4. Requests WITH the correct key should succeed")
    r = requests.post(f"{base_url}/add_project", json={"project": "test_correct_key"})
    r2 = requests.post(
        f"{base_url}/chat",
        json={"project": "test_correct_key", "message": "Say hello in one word."},
        headers={"X-API-Key": real_key},
    )
    check("/chat accepts the correct key", r2.status_code == 200, f"got {r2.status_code}: {r2.text[:200]}")


def test_rate_limit(base_url: str, real_key: str, rate_limit: int):
    print(f"\n5. Rate limiting should trigger after {rate_limit} requests/minute")
    hit_429 = False
    for i in range(rate_limit + 3):
        r = requests.post(
            f"{base_url}/search_web",
            json={"query": "test"},
            headers={"X-API-Key": real_key},
        )
        if r.status_code == 429:
            hit_429 = True
            check("Rate limit triggered", True, f"hit 429 on request #{i + 1}")
            break
    if not hit_429:
        check("Rate limit triggered", False, f"never got a 429 after {rate_limit + 3} requests")


def test_open_routes(base_url: str):
    print("\n6. Checking currently-unprotected routes (informational, not pass/fail)")
    r = requests.post(f"{base_url}/cancel")
    if r.status_code != 401:
        print(f"  [{WARN}] /cancel has no auth protection (got {r.status_code}) — decide if this needs one")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--url", default="http://localhost:8080", help="Base URL of the running app")
    parser.add_argument("--key", default=None, help="Your real API_KEY (or set it in .env)")
    parser.add_argument("--rate-limit", type=int, default=10, help="Expected RATE_LIMIT value from app.py")
    args = parser.parse_args()

    real_key = args.key or os.getenv("API_KEY")
    if not real_key:
        raise SystemExit("No API key provided. Use --key YOUR_KEY or set API_KEY in a .env file.")

    print(f"=== Testing {args.url} ===")
    test_no_key(args.url)
    test_wrong_key(args.url, real_key)
    test_substring_bug(args.url, real_key)
    test_correct_key(args.url, real_key)
    test_rate_limit(args.url, real_key, args.rate_limit)
    test_open_routes(args.url)
    print("\nDone.")