263 lines
8.2 KiB
Python
263 lines
8.2 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
|
|
from fastapi.testclient import TestClient
|
|
|
|
from device.manager import DeviceManager
|
|
from host_agent.config import HostAgentConfig
|
|
from host_agent.history import ConsoleHistoryStore
|
|
from host_agent.identity import HostIdentityStore
|
|
from host_agent.local_account import LocalAccountStore
|
|
from host_agent.status import AgentStatusTracker
|
|
from host_agent.web.app import SESSION_COOKIE_NAME, create_console_app
|
|
from host_agent.web.auth import SessionManager
|
|
from storage.device_config import DeviceConfigStore
|
|
|
|
CSRF_PATTERN = re.compile(r'name="csrf_token" value="([^"]+)"')
|
|
|
|
|
|
def _build_client(tmp_path, *, create_account: bool = True) -> tuple[TestClient, dict]:
|
|
config = HostAgentConfig(
|
|
control_plane_url="https://control.example",
|
|
host_id="host-a",
|
|
token="secret",
|
|
enrollment_managed=False,
|
|
console_session_ttl_seconds=3600.0,
|
|
)
|
|
manager = DeviceManager()
|
|
config_store = DeviceConfigStore(tmp_path / "devices.sqlite3")
|
|
local_account_store = LocalAccountStore(tmp_path / "host_local_account.json")
|
|
if create_account:
|
|
local_account_store.create("operator", "correct horse battery staple")
|
|
identity_store = HostIdentityStore(tmp_path / "host_identity.json")
|
|
history_store = ConsoleHistoryStore(tmp_path / "history.sqlite3")
|
|
status_tracker = AgentStatusTracker()
|
|
session_manager = SessionManager(ttl_seconds=3600.0)
|
|
|
|
app = create_console_app(
|
|
config=config,
|
|
manager=manager,
|
|
config_store=config_store,
|
|
local_account_store=local_account_store,
|
|
identity_store=identity_store,
|
|
history_store=history_store,
|
|
status_tracker=status_tracker,
|
|
session_manager=session_manager,
|
|
enrollment_client=None,
|
|
)
|
|
client = TestClient(app)
|
|
context = {
|
|
"manager": manager,
|
|
"config_store": config_store,
|
|
"local_account_store": local_account_store,
|
|
"history_store": history_store,
|
|
"session_manager": session_manager,
|
|
}
|
|
return client, context
|
|
|
|
|
|
def _login(
|
|
client: TestClient,
|
|
*,
|
|
username: str = "operator",
|
|
password: str = "correct horse battery staple",
|
|
) -> str:
|
|
response = client.post("/login", data={"username": username, "password": password})
|
|
assert response.status_code == 200
|
|
match = CSRF_PATTERN.search(response.text)
|
|
assert match is not None
|
|
return match.group(1)
|
|
|
|
|
|
def test_unauthenticated_get_root_redirects_to_login(tmp_path) -> None:
|
|
client, _ = _build_client(tmp_path)
|
|
|
|
response = client.get("/", follow_redirects=False)
|
|
|
|
assert response.status_code == 303
|
|
assert response.headers["location"] == "/login"
|
|
|
|
|
|
def test_login_with_no_account_shows_setup_message_and_rejects_post(tmp_path) -> None:
|
|
client, context = _build_client(tmp_path, create_account=False)
|
|
|
|
get_response = client.get("/login")
|
|
assert "device-host-agent setup" in get_response.text
|
|
|
|
post_response = client.post(
|
|
"/login", data={"username": "operator", "password": "anything"}
|
|
)
|
|
|
|
assert "device-host-agent setup" in post_response.text
|
|
assert SESSION_COOKIE_NAME not in client.cookies
|
|
|
|
|
|
def test_login_with_wrong_password_fails_and_sets_no_cookie(tmp_path) -> None:
|
|
client, _ = _build_client(tmp_path)
|
|
|
|
response = client.post("/login", data={"username": "operator", "password": "wrong"})
|
|
|
|
assert response.status_code == 200
|
|
assert "Invalid username or password" in response.text
|
|
assert SESSION_COOKIE_NAME not in client.cookies
|
|
|
|
|
|
def test_login_with_correct_password_sets_cookie_and_dashboard_succeeds(
|
|
tmp_path,
|
|
) -> None:
|
|
client, _ = _build_client(tmp_path)
|
|
|
|
response = client.post(
|
|
"/login",
|
|
data={"username": "operator", "password": "correct horse battery staple"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert SESSION_COOKIE_NAME in client.cookies
|
|
assert "Status" in response.text
|
|
assert "control.example" in response.text
|
|
|
|
|
|
def test_mutating_post_without_csrf_token_is_rejected_and_makes_no_change(
|
|
tmp_path,
|
|
) -> None:
|
|
client, context = _build_client(tmp_path)
|
|
_login(client)
|
|
|
|
response = client.post(
|
|
"/devices/save",
|
|
data={
|
|
"device_id": "device-a",
|
|
"driver_type": "wda",
|
|
"name": "Lab iPhone",
|
|
"connection_info": "{}",
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 403
|
|
assert context["config_store"].get("device-a") is None
|
|
assert context["manager"].list_devices() == []
|
|
|
|
|
|
def test_add_device_appears_in_devices_page_and_manager(tmp_path) -> None:
|
|
client, context = _build_client(tmp_path)
|
|
csrf_token = _login(client)
|
|
|
|
response = client.post(
|
|
"/devices/save",
|
|
data={
|
|
"device_id": "device-a",
|
|
"driver_type": "wda",
|
|
"name": "Lab iPhone",
|
|
"connection_info": "{}",
|
|
"csrf_token": csrf_token,
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert "device-a" in response.text
|
|
assert context["config_store"].get("device-a") is not None
|
|
assert [device.id for device in context["manager"].list_devices()] == ["device-a"]
|
|
|
|
|
|
def test_remove_device_unregisters_from_manager(tmp_path) -> None:
|
|
client, context = _build_client(tmp_path)
|
|
csrf_token = _login(client)
|
|
client.post(
|
|
"/devices/save",
|
|
data={
|
|
"device_id": "device-a",
|
|
"driver_type": "wda",
|
|
"name": "Lab iPhone",
|
|
"connection_info": "{}",
|
|
"csrf_token": csrf_token,
|
|
},
|
|
)
|
|
assert context["config_store"].get("device-a") is not None
|
|
|
|
response = client.post(
|
|
"/devices/remove",
|
|
data={"device_id": "device-a", "csrf_token": csrf_token},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert context["config_store"].get("device-a") is None
|
|
assert context["manager"].list_devices() == []
|
|
|
|
|
|
def test_change_password_wrong_current_fails_old_password_still_works(tmp_path) -> None:
|
|
client, context = _build_client(tmp_path)
|
|
csrf_token = _login(client)
|
|
|
|
response = client.post(
|
|
"/account",
|
|
data={
|
|
"current_password": "wrong",
|
|
"new_password": "new password",
|
|
"confirm_password": "new password",
|
|
"csrf_token": csrf_token,
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
assert "incorrect" in response.text
|
|
account_store: LocalAccountStore = context["local_account_store"]
|
|
account = account_store.load()
|
|
assert account is not None
|
|
assert account_store.verify(account, "correct horse battery staple") is True
|
|
|
|
|
|
def test_change_password_correct_succeeds_old_password_no_longer_works(
|
|
tmp_path,
|
|
) -> None:
|
|
client, context = _build_client(tmp_path)
|
|
csrf_token = _login(client)
|
|
|
|
response = client.post(
|
|
"/account",
|
|
data={
|
|
"current_password": "correct horse battery staple",
|
|
"new_password": "new password",
|
|
"confirm_password": "new password",
|
|
"csrf_token": csrf_token,
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert "Password updated" in response.text
|
|
account_store: LocalAccountStore = context["local_account_store"]
|
|
account = account_store.load()
|
|
assert account is not None
|
|
assert account_store.verify(account, "new password") is True
|
|
assert account_store.verify(account, "correct horse battery staple") is False
|
|
|
|
|
|
def test_history_reflects_entries_written_beforehand(tmp_path) -> None:
|
|
client, context = _build_client(tmp_path)
|
|
context["history_store"].record_heartbeat(device_count=3)
|
|
_login(client)
|
|
|
|
response = client.get("/history")
|
|
|
|
assert response.status_code == 200
|
|
assert "heartbeat: 3 devices" in response.text
|
|
|
|
|
|
def test_logout_invalidates_session_so_subsequent_request_redirects_to_login(
|
|
tmp_path,
|
|
) -> None:
|
|
client, _ = _build_client(tmp_path)
|
|
csrf_token = _login(client)
|
|
|
|
logout_response = client.post(
|
|
"/logout", data={"csrf_token": csrf_token}, follow_redirects=False
|
|
)
|
|
assert logout_response.status_code == 303
|
|
assert logout_response.headers["location"] == "/login"
|
|
|
|
response = client.get("/", follow_redirects=False)
|
|
|
|
assert response.status_code == 303
|
|
assert response.headers["location"] == "/login"
|