67 lines
2.2 KiB
Python
67 lines
2.2 KiB
Python
"""Database helpers.
|
|
|
|
Allows pre-loading database for tests.
|
|
"""
|
|
|
|
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
import httpx
|
|
from sshecret.crypto import generate_private_key, generate_public_key_string
|
|
|
|
|
|
class DatabasePreloader:
|
|
"""Database preloader class."""
|
|
|
|
def __init__(self, admin_url: str, username: str, password: str) -> None:
|
|
"""Instantiate class to populate database."""
|
|
self.admin_url: str = admin_url
|
|
self.username: str = username
|
|
self.password: str = password
|
|
|
|
@contextmanager
|
|
def login(self) -> Iterator[httpx.Client]:
|
|
"""Login and yield client."""
|
|
login_client = httpx.Client(base_url=self.admin_url)
|
|
resp = login_client.post(
|
|
"api/v1/token",
|
|
data={"username": self.username, "password": self.password}
|
|
)
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
token = data["access_token"]
|
|
headers = {"Authorization": f"Bearer {token}"}
|
|
with httpx.Client(base_url=self.admin_url, headers=headers) as client:
|
|
yield client
|
|
|
|
def create_client(self, *names: str) -> None:
|
|
"""Create one or more clients."""
|
|
with self.login() as http_client:
|
|
for name in names:
|
|
private_key = generate_private_key()
|
|
public_key = generate_public_key_string(private_key.public_key())
|
|
data = {
|
|
"name": name,
|
|
"description": "Test client",
|
|
"public_key": public_key,
|
|
"sources": ["0.0.0.0/0", "::/0"],
|
|
}
|
|
resp = http_client.post("api/v1/clients/", json=data)
|
|
resp.raise_for_status()
|
|
|
|
def create_secret(self, *secrets: tuple[str, list[str]]) -> None:
|
|
"""Create secret.
|
|
|
|
Argument format is (secret_name, [client1, client2, ...])
|
|
|
|
Clients must exist.
|
|
"""
|
|
with self.login() as http_client:
|
|
for name, clients in secrets:
|
|
data = {
|
|
"name": name,
|
|
"clients": clients,
|
|
"value": {"auto_generate": True, "length": 32}
|
|
}
|
|
|
|
http_client.post("api/v1/secrets/", json=data)
|