a789495a98
FreeBSD Smoke / FreeBSD Smoke (x86_64) (push) Has been cancelled
CI / Quality Guardrails (push) Has been cancelled
CI / Build & Test (macos-latest) (push) Has been cancelled
CI / Build & Test (ubuntu-latest) (push) Has been cancelled
CI / Build & Test (windows-latest) (push) Has been cancelled
CI / Format (push) Has been cancelled
CI / PowerShell Syntax (push) Has been cancelled
CI / Windows Cross-Target Check (Linux) (push) Has been cancelled
99 lines
3.3 KiB
Python
99 lines
3.3 KiB
Python
import importlib.util
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
os.environ["WAKE_TOKEN"] = "test-token"
|
|
os.environ["JCODE_GATEWAY_HOST"] = "100.64.0.10"
|
|
sys.modules.setdefault("boto3", types.ModuleType("boto3"))
|
|
sys.modules["boto3"].client = lambda *_args, **_kwargs: None
|
|
|
|
SPEC = importlib.util.spec_from_file_location(
|
|
"wake_lambda", Path(__file__).with_name("wake-lambda.py")
|
|
)
|
|
wake = importlib.util.module_from_spec(SPEC)
|
|
SPEC.loader.exec_module(wake)
|
|
|
|
|
|
class FakeEc2:
|
|
def __init__(self, state="running"):
|
|
self.state = state
|
|
self.started = False
|
|
|
|
def describe_instances(self, **_kwargs):
|
|
return {"Reservations": [{"Instances": [{"State": {"Name": self.state}}]}]}
|
|
|
|
def start_instances(self, **_kwargs):
|
|
self.started = True
|
|
|
|
|
|
class InvocationDoesNotExist(Exception):
|
|
pass
|
|
|
|
|
|
class FakeSsm:
|
|
class exceptions:
|
|
InvocationDoesNotExist = InvocationDoesNotExist
|
|
|
|
def describe_instance_information(self, **_kwargs):
|
|
return {"InstanceInformationList": [{"PingStatus": "Online"}]}
|
|
|
|
def send_command(self, **_kwargs):
|
|
return {"Command": {"CommandId": "command-1"}}
|
|
|
|
def get_command_invocation(self, **_kwargs):
|
|
return {
|
|
"Status": "Success",
|
|
"StandardOutputContent": "",
|
|
"StandardErrorContent": "Pairing code: \x1b[1;37m123 456\x1b[0m",
|
|
}
|
|
|
|
|
|
class WakeLambdaTests(unittest.TestCase):
|
|
def event(self, method, *, token=None, body=None, query=None):
|
|
headers = {"Authorization": f"Bearer {token}"} if token else {}
|
|
return {
|
|
"requestContext": {"http": {"method": method}},
|
|
"headers": headers,
|
|
"queryStringParameters": query,
|
|
"body": json.dumps(body) if body is not None else None,
|
|
}
|
|
|
|
def test_landing_page_never_embeds_token(self):
|
|
result = wake.handler(self.event("GET"), None)
|
|
self.assertEqual(result["statusCode"], 200)
|
|
self.assertNotIn("test-token", result["body"])
|
|
self.assertIn("Content-Security-Policy", result["headers"])
|
|
|
|
def test_legacy_query_redirects_token_into_fragment(self):
|
|
result = wake.handler(self.event("GET", query={"t": "test-token"}), None)
|
|
self.assertEqual(result["statusCode"], 302)
|
|
self.assertEqual(result["headers"]["Location"], "/#t=test-token")
|
|
|
|
def test_post_requires_bearer_token(self):
|
|
result = wake.handler(self.event("POST", body={"action": "status"}), None)
|
|
self.assertEqual(result["statusCode"], 403)
|
|
|
|
def test_status_uses_ec2_and_ssm_health(self):
|
|
ec2, ssm = FakeEc2(), FakeSsm()
|
|
with patch.object(wake.boto3, "client", side_effect=[ec2, ssm]):
|
|
result = wake.handler(
|
|
self.event("POST", token="test-token", body={"action": "status"}), None
|
|
)
|
|
self.assertEqual(result["statusCode"], 200)
|
|
self.assertTrue(json.loads(result["body"])["healthy"])
|
|
|
|
def test_pairing_parses_ansi_output_and_returns_tailnet_host(self):
|
|
result = wake.fetch_pair_code(FakeSsm())
|
|
self.assertEqual(result["code"], "123456")
|
|
self.assertEqual(result["host"], "100.64.0.10")
|
|
self.assertEqual(result["uri"], "jcode://pair?host=100.64.0.10&port=7643&code=123456")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|