@samitouri / QOSami_Agents_Zero-1
Code Commits Issues Pulls Sessions CI/CD
main
main
py 153 lines 4.21 KB
Raw
1 from __future__ import annotations
2
3 import json
4 from types import SimpleNamespace
5
6 import pytest
7 from flask import Flask
8
9
10 @pytest.mark.asyncio
11 async def test_csrf_token_allows_normalized_active_tailscale_origin(monkeypatch):
12 import api.csrf_token as csrf_module
13 import api.tunnel_proxy as tunnel_proxy
14
15 handler = csrf_module.GetCsrfToken(Flask("test_csrf_tunnel_origins"), None)
16 request = SimpleNamespace(
17 headers={"Origin": "https://agent-zero.tailabc.ts.net"},
18 environ={},
19 referrer=None,
20 )
21
22 monkeypatch.setattr(csrf_module.login, "is_login_required", lambda: False)
23 monkeypatch.setattr(
24 csrf_module.dotenv,
25 "get_dotenv_value",
26 lambda key: "http://localhost:32080"
27 if key == csrf_module.ALLOWED_ORIGINS_KEY
28 else "",
29 )
30
31 async def fake_tunnel_process(input_data):
32 return {
33 "success": True,
34 "tunnel_url": "https://agent-zero.tailabc.ts.net/funnel-ready/",
35 "is_running": True,
36 }
37
38 monkeypatch.setattr(tunnel_proxy, "process", fake_tunnel_process)
39
40 origin_check = await handler.check_allowed_origin(request)
41
42 assert origin_check["ok"] is True
43 assert "https://agent-zero.tailabc.ts.net" in origin_check["allowed_origins"]
44
45
46 @pytest.mark.asyncio
47 async def test_csrf_token_rejects_unrelated_origin_with_active_tunnel(monkeypatch):
48 import api.csrf_token as csrf_module
49 import api.tunnel_proxy as tunnel_proxy
50
51 handler = csrf_module.GetCsrfToken(Flask("test_csrf_tunnel_origins"), None)
52 request = SimpleNamespace(
53 headers={"Origin": "https://evil.example"},
54 environ={},
55 referrer=None,
56 )
57
58 monkeypatch.setattr(csrf_module.login, "is_login_required", lambda: False)
59 monkeypatch.setattr(
60 csrf_module.dotenv,
61 "get_dotenv_value",
62 lambda key: "http://localhost:32080"
63 if key == csrf_module.ALLOWED_ORIGINS_KEY
64 else "",
65 )
66
67 async def fake_tunnel_process(input_data):
68 return {
69 "success": True,
70 "tunnel_url": "https://agent-zero.tailabc.ts.net/funnel-ready/",
71 "is_running": True,
72 }
73
74 monkeypatch.setattr(tunnel_proxy, "process", fake_tunnel_process)
75
76 origin_check = await handler.check_allowed_origin(request)
77
78 assert origin_check["ok"] is False
79
80
81 def test_active_tunnel_origins_include_docker_tunnel_service_url(monkeypatch):
82 import helpers.tunnel_origins as tunnel_origins
83
84 monkeypatch.setattr(
85 tunnel_origins,
86 "_get_tunnel_service_url",
87 lambda: "https://agent-zero.tailabc.ts.net/funnel-ready/",
88 )
89
90 assert (
91 "https://agent-zero.tailabc.ts.net"
92 in tunnel_origins.get_active_tunnel_origins()
93 )
94
95
96 def test_tunnel_service_url_uses_short_local_get_request(monkeypatch):
97 from helpers import dotenv, runtime
98 import helpers.tunnel_origins as tunnel_origins
99
100 captured = {}
101
102 class FakeResponse:
103 def __enter__(self):
104 return self
105
106 def __exit__(self, exc_type, exc, tb):
107 return None
108
109 def read(self):
110 return json.dumps({
111 "success": True,
112 "tunnel_url": "https://agent-zero.tailabc.ts.net/funnel-ready/",
113 }).encode("utf-8")
114
115 def fake_urlopen(request, timeout):
116 captured["url"] = request.full_url
117 captured["body"] = request.data
118 captured["method"] = request.get_method()
119 captured["timeout"] = timeout
120 return FakeResponse()
121
122 monkeypatch.setattr(
123 runtime,
124 "is_dockerized",
125 lambda: True,
126 )
127 monkeypatch.setattr(
128 runtime,
129 "get_arg",
130 lambda name: None,
131 )
132 monkeypatch.setattr(
133 runtime,
134 "get_tunnel_api_port",
135 lambda: 55520,
136 )
137 monkeypatch.setattr(
138 dotenv,
139 "get_dotenv_value",
140 lambda key: "",
141 )
142 monkeypatch.setattr(tunnel_origins.urllib.request, "urlopen", fake_urlopen)
143
144 assert (
145 tunnel_origins._get_tunnel_service_url()
146 == "https://agent-zero.tailabc.ts.net/funnel-ready/"
147 )
148 assert captured == {
149 "url": "http://localhost:55520/",
150 "body": b'{"action": "get"}',
151 "method": "POST",
152 "timeout": 0.35,
153 }