100 lines
3.9 KiB
Python
100 lines
3.9 KiB
Python
import hashlib
|
|
import hmac
|
|
import unittest
|
|
import urllib.error
|
|
from unittest.mock import patch
|
|
|
|
import webssh_backend.telemetry as telemetry
|
|
from webssh_backend.telemetry import TelemetryClient, canonical_payload, sign_request
|
|
|
|
|
|
class TelemetrySigningTests(unittest.TestCase):
|
|
def test_canonical_payload_has_exact_newlines_and_body_hash(self):
|
|
body = b'{"version":"1.0.0"}'
|
|
expected = "POST\n/api/v1/telemetry/heartbeat\n1700000000\nnonce\nrequest\n" + hashlib.sha256(body).hexdigest()
|
|
self.assertEqual(canonical_payload("post", "/api/v1/telemetry/heartbeat", 1700000000, "nonce", "request", body), expected)
|
|
self.assertFalse(expected.endswith("\n"))
|
|
|
|
def test_signature_is_lowercase_hmac_sha256(self):
|
|
body = b"{}"
|
|
canonical = canonical_payload("POST", "/path", 1, "n", "r", body)
|
|
expected = hmac.new(b"secret", canonical.encode(), hashlib.sha256).hexdigest()
|
|
self.assertEqual(sign_request("secret", "POST", "/path", 1, "n", "r", body), expected)
|
|
self.assertEqual(len(expected), 64)
|
|
self.assertEqual(expected, expected.lower())
|
|
|
|
|
|
class FakeResponse:
|
|
status = 201
|
|
headers = {}
|
|
|
|
def __init__(self, payload):
|
|
self.payload = payload
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
return False
|
|
|
|
def read(self):
|
|
import json
|
|
return json.dumps(self.payload).encode()
|
|
|
|
|
|
class FakeState:
|
|
def __init__(self, state=None):
|
|
self.state = state or {}
|
|
|
|
def get_telemetry_state(self):
|
|
return dict(self.state) if self.state else None
|
|
|
|
def save_telemetry_state(self, state):
|
|
self.state = dict(state)
|
|
return True
|
|
|
|
|
|
class TelemetryClientTests(unittest.TestCase):
|
|
def test_disabled_telemetry_does_not_call_opener(self):
|
|
called = []
|
|
|
|
def opener(request, timeout):
|
|
called.append(request)
|
|
return FakeResponse({"data": {}})
|
|
|
|
with patch.object(telemetry, "creds_manager", FakeState({"credential_secret": "secret"})), patch.object(telemetry, "STRIX_TELEMETRY_ENABLED", False), patch.object(telemetry, "STRIX_TELEMETRY_URL", "https://telemetry.test"):
|
|
TelemetryClient(opener).uninstall()
|
|
self.assertEqual(called, [])
|
|
|
|
def test_registration_persists_installation_and_credential(self):
|
|
state = FakeState()
|
|
opener = lambda request, timeout: FakeResponse({"data": {
|
|
"installation_uuid": state.state["installation_uuid"],
|
|
"credential": "credential-id.secret-value",
|
|
"expires_at": "2099-01-01T00:00:00Z",
|
|
}})
|
|
with patch.object(telemetry, "creds_manager", state), patch.object(telemetry, "STRIX_TELEMETRY_URL", "https://telemetry.test"):
|
|
TelemetryClient(opener).register()
|
|
self.assertTrue(state.state["installation_uuid"])
|
|
self.assertEqual(state.state["credential_id"], "credential-id")
|
|
self.assertEqual(state.state["credential_secret"], "secret-value")
|
|
|
|
def test_retry_reuses_request_id_and_regenerates_nonce(self):
|
|
state = FakeState({
|
|
"installation_uuid": "installation",
|
|
"credential_id": "credential-id",
|
|
"credential_secret": "secret-value",
|
|
})
|
|
requests = []
|
|
|
|
def opener(request, timeout):
|
|
requests.append(request)
|
|
if len(requests) == 1:
|
|
raise urllib.error.URLError("temporary failure")
|
|
return FakeResponse({"data": {"accepted": True}})
|
|
|
|
with patch.object(telemetry, "creds_manager", state), patch.object(telemetry, "STRIX_TELEMETRY_URL", "https://telemetry.test"), patch.object(telemetry.time, "sleep"):
|
|
TelemetryClient(opener).heartbeat()
|
|
self.assertEqual(len(requests), 2)
|
|
self.assertEqual(requests[0].headers["X-strix-request-id"], requests[1].headers["X-strix-request-id"])
|
|
self.assertNotEqual(requests[0].headers["X-strix-nonce"], requests[1].headers["X-strix-nonce"]) |