feat(cloud): add edge host enrollment
This commit is contained in:
@@ -7,10 +7,12 @@ from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import Engine, delete, func, select
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from cloud.db_models import (
|
||||
Base,
|
||||
DeviceEnrollmentRow,
|
||||
HostRow,
|
||||
PluginRow,
|
||||
PooledDeviceRow,
|
||||
@@ -33,6 +35,211 @@ class SQLAlchemyCloudRepository:
|
||||
if create_schema:
|
||||
Base.metadata.create_all(engine)
|
||||
|
||||
def enroll_host(
|
||||
self,
|
||||
*,
|
||||
host_id: str,
|
||||
agent_instance_id: str,
|
||||
credential_digest: str,
|
||||
enrollment_token_digest: str,
|
||||
display_name: str | None,
|
||||
enrolled_at: datetime,
|
||||
) -> Any:
|
||||
from cloud.repository import (
|
||||
EnrollmentTokenConflictError,
|
||||
HostEnrollmentConflictError,
|
||||
)
|
||||
|
||||
try:
|
||||
with self._sessions.begin() as session:
|
||||
statement = select(HostRow).where(
|
||||
HostRow.agent_instance_id == agent_instance_id
|
||||
)
|
||||
if self.engine.dialect.name == "postgresql":
|
||||
statement = statement.with_for_update()
|
||||
existing = session.scalars(statement).first()
|
||||
if existing is not None:
|
||||
if (
|
||||
existing.credential_digest != credential_digest
|
||||
or existing.enrollment_token_digest != enrollment_token_digest
|
||||
):
|
||||
raise HostEnrollmentConflictError(
|
||||
"Host enrollment identity does not match existing binding"
|
||||
)
|
||||
return _host_enrollment_from_row(existing)
|
||||
|
||||
token_owner = session.scalars(
|
||||
select(HostRow).where(
|
||||
HostRow.enrollment_token_digest == enrollment_token_digest
|
||||
)
|
||||
).first()
|
||||
if token_owner is not None:
|
||||
raise EnrollmentTokenConflictError(
|
||||
"enrollment token is already bound to another Host"
|
||||
)
|
||||
|
||||
row = HostRow(
|
||||
host_id=host_id,
|
||||
address=None,
|
||||
last_seen_at=_iso(enrolled_at),
|
||||
agent_instance_id=agent_instance_id,
|
||||
credential_digest=credential_digest,
|
||||
enrollment_token_digest=enrollment_token_digest,
|
||||
display_name=display_name,
|
||||
enrolled_at=_iso(enrolled_at),
|
||||
revoked_at=None,
|
||||
)
|
||||
session.add(row)
|
||||
session.flush()
|
||||
return _host_enrollment_from_row(row)
|
||||
except IntegrityError as exc:
|
||||
with self._sessions() as session:
|
||||
existing = session.scalars(
|
||||
select(HostRow).where(
|
||||
HostRow.agent_instance_id == agent_instance_id
|
||||
)
|
||||
).first()
|
||||
if (
|
||||
existing is not None
|
||||
and existing.credential_digest == credential_digest
|
||||
and existing.enrollment_token_digest == enrollment_token_digest
|
||||
):
|
||||
return _host_enrollment_from_row(existing)
|
||||
token_owner = session.scalars(
|
||||
select(HostRow).where(
|
||||
HostRow.enrollment_token_digest == enrollment_token_digest
|
||||
)
|
||||
).first()
|
||||
if token_owner is not None:
|
||||
raise EnrollmentTokenConflictError(
|
||||
"enrollment token is already bound to another Host"
|
||||
) from exc
|
||||
raise HostEnrollmentConflictError(
|
||||
"Host enrollment conflicts with an existing identity"
|
||||
) from exc
|
||||
|
||||
def authenticate_enrolled_host(self, credential_digest: str) -> str | None:
|
||||
with self._sessions() as session:
|
||||
host_id = session.scalar(
|
||||
select(HostRow.host_id)
|
||||
.where(
|
||||
HostRow.credential_digest == credential_digest,
|
||||
HostRow.revoked_at.is_(None),
|
||||
)
|
||||
.limit(1)
|
||||
)
|
||||
return str(host_id) if host_id is not None else None
|
||||
|
||||
def revoke_enrolled_host(
|
||||
self,
|
||||
host_id: str,
|
||||
*,
|
||||
revoked_at: datetime,
|
||||
) -> bool:
|
||||
with self._sessions.begin() as session:
|
||||
row = session.get(
|
||||
HostRow,
|
||||
host_id,
|
||||
with_for_update=self.engine.dialect.name == "postgresql",
|
||||
)
|
||||
if row is None or row.credential_digest is None:
|
||||
return False
|
||||
row.revoked_at = _iso(revoked_at)
|
||||
return True
|
||||
|
||||
def is_enrollment_managed_host(self, host_id: str) -> bool:
|
||||
with self._sessions() as session:
|
||||
credential_digest = session.scalar(
|
||||
select(HostRow.credential_digest).where(HostRow.host_id == host_id)
|
||||
)
|
||||
return credential_digest is not None
|
||||
|
||||
def enroll_device(
|
||||
self,
|
||||
*,
|
||||
device_id: str,
|
||||
host_id: str,
|
||||
local_device_id: str,
|
||||
driver_type: str,
|
||||
name: str | None,
|
||||
capability_tags: list[str],
|
||||
enrolled_at: datetime,
|
||||
) -> Any:
|
||||
from cloud.repository import DeviceEnrollmentConflictError
|
||||
|
||||
tags_json = json.dumps(list(capability_tags), ensure_ascii=False)
|
||||
try:
|
||||
with self._sessions.begin() as session:
|
||||
statement = select(DeviceEnrollmentRow).where(
|
||||
DeviceEnrollmentRow.host_id == host_id,
|
||||
DeviceEnrollmentRow.local_device_id == local_device_id,
|
||||
)
|
||||
if self.engine.dialect.name == "postgresql":
|
||||
statement = statement.with_for_update()
|
||||
existing = session.scalars(statement).first()
|
||||
if existing is not None:
|
||||
if existing.driver_type != driver_type:
|
||||
raise DeviceEnrollmentConflictError(
|
||||
"device driver type does not match existing enrollment"
|
||||
)
|
||||
if existing.revoked_at is not None:
|
||||
raise DeviceEnrollmentConflictError(
|
||||
"device enrollment has been revoked"
|
||||
)
|
||||
existing.name = name
|
||||
existing.capability_tags_json = tags_json
|
||||
return _device_enrollment_from_row(existing)
|
||||
|
||||
host = session.get(HostRow, host_id)
|
||||
if host is None:
|
||||
session.add(
|
||||
HostRow(
|
||||
host_id=host_id,
|
||||
address=None,
|
||||
last_seen_at=_iso(enrolled_at),
|
||||
)
|
||||
)
|
||||
row = DeviceEnrollmentRow(
|
||||
device_id=device_id,
|
||||
host_id=host_id,
|
||||
local_device_id=local_device_id,
|
||||
driver_type=driver_type,
|
||||
name=name,
|
||||
capability_tags_json=tags_json,
|
||||
enrolled_at=_iso(enrolled_at),
|
||||
revoked_at=None,
|
||||
)
|
||||
session.add(row)
|
||||
session.flush()
|
||||
return _device_enrollment_from_row(row)
|
||||
except IntegrityError as exc:
|
||||
with self._sessions() as session:
|
||||
existing = session.scalars(
|
||||
select(DeviceEnrollmentRow).where(
|
||||
DeviceEnrollmentRow.host_id == host_id,
|
||||
DeviceEnrollmentRow.local_device_id == local_device_id,
|
||||
)
|
||||
).first()
|
||||
if existing is not None and existing.driver_type == driver_type:
|
||||
return _device_enrollment_from_row(existing)
|
||||
raise DeviceEnrollmentConflictError(
|
||||
"device enrollment conflicts with an existing identity"
|
||||
) from exc
|
||||
|
||||
def get_device_enrollment(self, device_id: str) -> Any | None:
|
||||
with self._sessions() as session:
|
||||
row = session.get(DeviceEnrollmentRow, device_id)
|
||||
return _device_enrollment_from_row(row) if row else None
|
||||
|
||||
def list_device_enrollments(self, host_id: str) -> list[Any]:
|
||||
with self._sessions() as session:
|
||||
rows = session.scalars(
|
||||
select(DeviceEnrollmentRow)
|
||||
.where(DeviceEnrollmentRow.host_id == host_id)
|
||||
.order_by(DeviceEnrollmentRow.device_id)
|
||||
).all()
|
||||
return [_device_enrollment_from_row(row) for row in rows]
|
||||
|
||||
def upsert_host(
|
||||
self,
|
||||
host_id: str,
|
||||
@@ -568,6 +775,40 @@ def _host_from_row(row: HostRow) -> Any:
|
||||
)
|
||||
|
||||
|
||||
def _host_enrollment_from_row(row: HostRow) -> Any:
|
||||
from cloud.repository import HostEnrollment
|
||||
|
||||
enrolled_at = _parse_dt(row.enrolled_at) or _parse_dt(row.last_seen_at) or utc_now()
|
||||
if row.agent_instance_id is None:
|
||||
raise ValueError(f"Host {row.host_id!r} is not enrollment-managed")
|
||||
return HostEnrollment(
|
||||
host_id=row.host_id,
|
||||
agent_instance_id=row.agent_instance_id,
|
||||
display_name=row.display_name,
|
||||
enrolled_at=enrolled_at,
|
||||
revoked_at=_parse_dt(row.revoked_at),
|
||||
)
|
||||
|
||||
|
||||
def _device_enrollment_from_row(row: DeviceEnrollmentRow) -> Any:
|
||||
from cloud.repository import DeviceEnrollment
|
||||
|
||||
try:
|
||||
tags = list(json.loads(row.capability_tags_json))
|
||||
except TypeError, ValueError:
|
||||
tags = []
|
||||
return DeviceEnrollment(
|
||||
device_id=row.device_id,
|
||||
host_id=row.host_id,
|
||||
local_device_id=row.local_device_id,
|
||||
driver_type=row.driver_type,
|
||||
name=row.name,
|
||||
capability_tags=tags,
|
||||
enrolled_at=_parse_dt(row.enrolled_at) or utc_now(),
|
||||
revoked_at=_parse_dt(row.revoked_at),
|
||||
)
|
||||
|
||||
|
||||
def _device_from_row(row: PooledDeviceRow) -> Any:
|
||||
from cloud.pool import PooledDevice
|
||||
|
||||
|
||||
Reference in New Issue
Block a user