UncountablePythonSDK 0.0.8__py3-none-any.whl → 0.0.92__py3-none-any.whl

uncountable/integration/server.py

@@ -0,0 +1,117 @@
+import signal
+from dataclasses import asdict
+from types import TracebackType
+from typing import Optional, assert_never
+
+from apscheduler.executors.pool import ThreadPoolExecutor
+from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
+from apscheduler.schedulers.background import BackgroundScheduler
+from apscheduler.schedulers.base import BaseScheduler
+from apscheduler.triggers.cron import CronTrigger
+from opentelemetry.trace import get_current_span
+from sqlalchemy.engine.base import Engine
+
+from uncountable.integration.cron import CronJobArgs, cron_job_executor
+from uncountable.integration.telemetry import Logger
+from uncountable.types import base_t, job_definition_t
+from uncountable.types.job_definition_t import (
+    CronJobDefinition,
+    WebhookJobDefinition,
+)
+
+_MAX_APSCHEDULER_CONCURRENT_JOBS = 1
+
+
+class IntegrationServer:
+    _scheduler: BaseScheduler
+    _engine: Engine
+    _server_logger: Logger
+
+    def __init__(self, engine: Engine) -> None:
+        self._engine = engine
+        self._scheduler = BackgroundScheduler(
+            timezone="UTC",
+            jobstores={"default": SQLAlchemyJobStore(engine=engine)},
+            executors={"default": ThreadPoolExecutor(_MAX_APSCHEDULER_CONCURRENT_JOBS)},
+        )
+        self._server_logger = Logger(get_current_span())
+
+    def register_jobs(self, profiles: list[job_definition_t.ProfileMetadata]) -> None:
+        valid_job_ids = []
+        for profile_metadata in profiles:
+            for job_defn in profile_metadata.jobs:
+                valid_job_ids.append(job_defn.id)
+                match job_defn:
+                    case CronJobDefinition():
+                        # Add to ap scheduler
+                        job_kwargs = asdict(
+                            CronJobArgs(
+                                definition=job_defn, profile_metadata=profile_metadata
+                            )
+                        )
+                        try:
+                            existing_job = self._scheduler.get_job(job_defn.id)
+                        except ValueError as e:
+                            self._server_logger.log_warning(
+                                f"could not reconstitute job {job_defn.id}: {e}"
+                            )
+                            self._scheduler.remove_job(job_defn.id)
+                            existing_job = None
+                        if existing_job is not None:
+                            existing_job.modify(
+                                name=job_defn.name,
+                                kwargs=job_kwargs,
+                                misfire_grace_time=None,
+                            )
+                            existing_job.reschedule(
+                                CronTrigger.from_crontab(job_defn.cron_spec)
+                            )
+                            if not job_defn.enabled:
+                                existing_job.pause()
+                            else:
+                                existing_job.resume()
+                        else:
+                            job_opts: dict[str, base_t.JsonValue] = {}
+                            if not job_defn.enabled:
+                                job_opts["next_run_time"] = None
+                            self._scheduler.add_job(
+                                cron_job_executor,
+                                # IMPROVE: reconsider these defaults
+                                max_instances=1,
+                                coalesce=True,
+                                trigger=CronTrigger.from_crontab(job_defn.cron_spec),
+                                name=job_defn.name,
+                                id=job_defn.id,
+                                kwargs=job_kwargs,
+                                misfire_grace_time=None,
+                                **job_opts,
+                            )
+                    case WebhookJobDefinition():
+                        pass
+                    case _:
+                        assert_never(job_defn)
+        all_jobs = self._scheduler.get_jobs()
+        for job in all_jobs:
+            if job.id not in valid_job_ids:
+                self._scheduler.remove_job(job.id)
+
+    def serve_forever(self) -> None:
+        signal.pause()
+
+    def _start_apscheduler(self) -> None:
+        self._scheduler.start()
+
+    def _stop_apscheduler(self) -> None:
+        self._scheduler.shutdown()
+
+    def __enter__(self) -> "IntegrationServer":
+        self._start_apscheduler()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        self._stop_apscheduler()

uncountable/integration/telemetry.py

@@ -0,0 +1,209 @@
+import functools
+import os
+import time
+import traceback
+import typing
+from contextlib import contextmanager
+from enum import StrEnum
+from typing import Generator, assert_never, cast
+
+from opentelemetry import _logs, trace
+from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk._logs import Logger as OTELLogger
+from opentelemetry.sdk._logs import LoggerProvider, LogRecord
+from opentelemetry.sdk._logs.export import BatchLogRecordProcessor, ConsoleLogExporter
+from opentelemetry.sdk.resources import Attributes, Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import (
+    SimpleSpanProcessor,
+)
+from opentelemetry.trace import DEFAULT_TRACE_OPTIONS, Span, Tracer
+
+from uncountable.core.environment import (
+    get_otel_enabled,
+    get_server_env,
+    get_version,
+)
+from uncountable.types import base_t, job_definition_t
+
+
+def _cast_attributes(attributes: dict[str, base_t.JsonValue]) -> Attributes:
+    return cast(Attributes, attributes)
+
+
+@functools.cache
+def get_otel_resource() -> Resource:
+    attributes: dict[str, base_t.JsonValue] = {
+        "service.name": "integration-server",
+        "sdk.version": get_version(),
+    }
+    unc_version = os.environ.get("UNC_VERSION")
+    if unc_version is not None:
+        attributes["service.version"] = unc_version
+    unc_env = get_server_env()
+    if unc_env is not None:
+        attributes["deployment.environment"] = unc_env
+    resource = Resource.create(attributes=_cast_attributes(attributes))
+    return resource
+
+
+@functools.cache
+def get_otel_tracer() -> Tracer:
+    provider = TracerProvider(resource=get_otel_resource())
+    if get_otel_enabled():
+        provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter()))
+    trace.set_tracer_provider(provider)
+    return provider.get_tracer("integration.telemetry")
+
+
+@functools.cache
+def get_otel_logger() -> OTELLogger:
+    provider = LoggerProvider(resource=get_otel_resource())
+    provider.add_log_record_processor(BatchLogRecordProcessor(ConsoleLogExporter()))
+    if get_otel_enabled():
+        provider.add_log_record_processor(BatchLogRecordProcessor(OTLPLogExporter()))
+    _logs.set_logger_provider(provider)
+    return provider.get_logger("integration.telemetry")
+
+
+class LogSeverity(StrEnum):
+    INFO = "Info"
+    WARN = "Warn"
+    ERROR = "Error"
+
+
+class Logger:
+    current_span: Span
+
+    def __init__(self, base_span: Span) -> None:
+        self.current_span = base_span
+
+    @property
+    def current_span_id(self) -> int:
+        return self.current_span.get_span_context().span_id
+
+    @property
+    def current_trace_id(self) -> int | None:
+        return self.current_span.get_span_context().trace_id
+
+    def _patch_attributes(self, attributes: Attributes | None) -> Attributes:
+        return attributes or {}
+
+    def _emit_log(
+        self, message: str, *, severity: LogSeverity, attributes: Attributes | None
+    ) -> None:
+        otel_logger = get_otel_logger()
+        log_record = LogRecord(
+            body=message,
+            severity_text=severity,
+            timestamp=time.time_ns(),
+            attributes=self._patch_attributes(attributes),
+            span_id=self.current_span_id,
+            trace_id=self.current_trace_id,
+            trace_flags=DEFAULT_TRACE_OPTIONS,
+            severity_number=_logs.SeverityNumber.UNSPECIFIED,
+            resource=get_otel_resource(),
+        )
+        otel_logger.emit(log_record)
+
+    def log_info(self, message: str, *, attributes: Attributes | None = None) -> None:
+        self._emit_log(
+            message=message, severity=LogSeverity.INFO, attributes=attributes
+        )
+
+    def log_warning(
+        self, message: str, *, attributes: Attributes | None = None
+    ) -> None:
+        self._emit_log(
+            message=message, severity=LogSeverity.WARN, attributes=attributes
+        )
+
+    def log_error(self, message: str, *, attributes: Attributes | None = None) -> None:
+        self._emit_log(
+            message=message, severity=LogSeverity.ERROR, attributes=attributes
+        )
+
+    def log_exception(
+        self,
+        exception: BaseException,
+        *,
+        message: str | None = None,
+        attributes: Attributes | None = None,
+    ) -> None:
+        traceback_str = "".join(traceback.format_tb(exception.__traceback__))
+        patched_attributes = self._patch_attributes(attributes)
+        self.current_span.record_exception(
+            exception=exception, attributes=patched_attributes
+        )
+        self.log_error(
+            message=f"error: {message}\nexception: {exception}{traceback_str}",
+            attributes=patched_attributes,
+        )
+
+    @contextmanager
+    def push_scope(
+        self, scope_name: str, *, attributes: Attributes | None = None
+    ) -> Generator[typing.Self, None, None]:
+        with get_otel_tracer().start_as_current_span(
+            scope_name, attributes=self._patch_attributes(attributes)
+        ):
+            yield self
+
+
+class JobLogger(Logger):
+    def __init__(
+        self,
+        *,
+        base_span: Span,
+        profile_metadata: job_definition_t.ProfileMetadata,
+        job_definition: job_definition_t.JobDefinition,
+    ) -> None:
+        self.profile_metadata = profile_metadata
+        self.job_definition = job_definition
+        super().__init__(base_span)
+
+    def _patch_attributes(self, attributes: Attributes | None) -> Attributes:
+        patched_attributes: dict[str, base_t.JsonValue] = {
+            **(attributes if attributes is not None else {})
+        }
+        patched_attributes["profile.name"] = self.profile_metadata.name
+        patched_attributes["profile.base_url"] = self.profile_metadata.base_url
+        patched_attributes["job.name"] = self.job_definition.name
+        patched_attributes["job.id"] = self.job_definition.id
+        patched_attributes["job.definition_type"] = self.job_definition.type
+        match self.job_definition:
+            case job_definition_t.CronJobDefinition():
+                patched_attributes["job.definition.cron_spec"] = (
+                    self.job_definition.cron_spec
+                )
+            case job_definition_t.WebhookJobDefinition():
+                pass
+            case _:
+                assert_never(self.job_definition)
+        patched_attributes["job.definition.executor.type"] = (
+            self.job_definition.executor.type
+        )
+        match self.job_definition.executor:
+            case job_definition_t.JobExecutorScript():
+                patched_attributes["job.definition.executor.import_path"] = (
+                    self.job_definition.executor.import_path
+                )
+            case job_definition_t.JobExecutorGenericUpload():
+                patched_attributes["job.definition.executor.data_source.type"] = (
+                    self.job_definition.executor.data_source.type
+                )
+            case _:
+                assert_never(self.job_definition.executor)
+        return _cast_attributes(patched_attributes)
+
+
+@contextmanager
+def push_scope_optional(
+    logger: Logger | None, scope_name: str, *, attributes: Attributes | None = None
+) -> Generator[None, None, None]:
+    if logger is None:
+        yield
+    else:
+        with logger.push_scope(scope_name, attributes=attributes):
+            yield

uncountable/integration/webhook_server/entrypoint.py

@@ -0,0 +1,170 @@
+import hmac
+import typing
+from dataclasses import dataclass
+
+import flask
+import simplejson
+from flask.typing import ResponseReturnValue
+from flask.wrappers import Response
+from opentelemetry.trace import get_current_span
+from uncountable.core.environment import (
+    get_local_admin_server_port,
+    get_server_env,
+    get_webhook_server_port,
+)
+from uncountable.integration.queue_runner.command_server.command_client import (
+    send_job_queue_message,
+)
+from uncountable.integration.queue_runner.command_server.types import (
+    CommandServerException,
+)
+from uncountable.integration.scan_profiles import load_profiles
+from uncountable.integration.secret_retrieval.retrieve_secret import retrieve_secret
+from uncountable.integration.telemetry import Logger
+from uncountable.types import base_t, job_definition_t, queued_job_t, webhook_job_t
+
+from pkgs.argument_parser import CachedParser
+
+app = flask.Flask(__name__)
+
+
+@dataclass(kw_only=True)
+class WebhookResponse:
+    pass
+
+
+webhook_payload_parser = CachedParser(webhook_job_t.WebhookEventBody)
+
+
+class WebhookException(BaseException):
+    error_code: int
+    message: str
+
+    def __init__(self, *, error_code: int, message: str) -> None:
+        self.error_code = error_code
+        self.message = message
+
+    @staticmethod
+    def payload_failed_signature() -> "WebhookException":
+        return WebhookException(
+            error_code=401, message="webhook payload did not match signature"
+        )
+
+    @staticmethod
+    def no_signature_passed() -> "WebhookException":
+        return WebhookException(error_code=400, message="missing signature")
+
+    @staticmethod
+    def body_parse_error() -> "WebhookException":
+        return WebhookException(error_code=400, message="body parse error")
+
+    @staticmethod
+    def unknown_error() -> "WebhookException":
+        return WebhookException(error_code=500, message="internal server error")
+
+    def __str__(self) -> str:
+        return f"[{self.error_code}]: {self.message}"
+
+    def make_error_response(self) -> Response:
+        return Response(
+            status=self.error_code, response={"error": {"message": str(self)}}
+        )
+
+
+def _parse_webhook_payload(
+    *, raw_request_body: bytes, signature_key: str, passed_signature: str
+) -> base_t.JsonValue:
+    request_body_signature = hmac.new(
+        signature_key.encode("utf-8"), msg=raw_request_body, digestmod="sha256"
+    ).hexdigest()
+
+    if request_body_signature != passed_signature:
+        raise WebhookException.payload_failed_signature()
+
+    try:
+        request_body = simplejson.loads(raw_request_body.decode())
+        return typing.cast(base_t.JsonValue, request_body)
+    except (simplejson.JSONDecodeError, ValueError) as e:
+        raise WebhookException.body_parse_error() from e
+
+
+def register_route(
+    *,
+    server_logger: Logger,
+    profile_meta: job_definition_t.ProfileMetadata,
+    job: job_definition_t.WebhookJobDefinition,
+) -> None:
+    route = f"/{profile_meta.name}/{job.id}"
+
+    def handle_webhook() -> ResponseReturnValue:
+        with server_logger.push_scope(route):
+            try:
+                signature_key = retrieve_secret(
+                    profile_metadata=profile_meta,
+                    secret_retrieval=job.signature_key_secret,
+                )
+
+                passed_signature = flask.request.headers.get(
+                    "Uncountable-Webhook-Signature"
+                )
+                if passed_signature is None:
+                    raise WebhookException.no_signature_passed()
+
+                webhook_payload = _parse_webhook_payload(
+                    raw_request_body=flask.request.data,
+                    signature_key=signature_key,
+                    passed_signature=passed_signature,
+                )
+
+                try:
+                    send_job_queue_message(
+                        job_ref_name=job.id,
+                        payload=queued_job_t.QueuedJobPayload(
+                            invocation_context=queued_job_t.InvocationContextWebhook(
+                                webhook_payload=webhook_payload
+                            )
+                        ),
+                        port=get_local_admin_server_port(),
+                    )
+                except CommandServerException as e:
+                    raise WebhookException.unknown_error() from e
+
+                return flask.jsonify(WebhookResponse())
+            except WebhookException as e:
+                server_logger.log_exception(e)
+                return e.make_error_response()
+            except Exception as e:
+                server_logger.log_exception(e)
+                return WebhookException.unknown_error().make_error_response()
+
+    app.add_url_rule(
+        route,
+        endpoint=f"handle_webhook_{job.id}",
+        view_func=handle_webhook,
+        methods=["POST"],
+    )
+
+    server_logger.log_info(f"job {job.id} webhook registered at: {route}")
+
+
+def main() -> None:
+    profiles = load_profiles()
+    for profile_metadata in profiles:
+        server_logger = Logger(get_current_span())
+        for job in profile_metadata.jobs:
+            if isinstance(job, job_definition_t.WebhookJobDefinition):
+                register_route(
+                    server_logger=server_logger, profile_meta=profile_metadata, job=job
+                )
+
+
+main()
+
+
+if __name__ == "__main__":
+    app.run(
+        host="0.0.0.0",
+        port=get_webhook_server_port(),
+        debug=get_server_env() == "playground",
+        exclude_patterns=[],
+    )