ReticulumTelemetryHub 0.1.0__py3-none-any.whl → 0.143.0__py3-none-any.whl

reticulum_telemetry_hub/atak_cot/pytak_client.py

@@ -0,0 +1,569 @@
+"""PyTAK client helpers for sending and receiving Cursor on Target events."""
+
+from __future__ import annotations
+
+import asyncio
+import atexit
+import logging
+import sys
+import types
+import weakref
+import xml.etree.ElementTree as ET
+from configparser import ConfigParser, SectionProxy
+from contextlib import suppress
+from importlib.util import find_spec
+from threading import Event as ThreadEvent
+from threading import Lock
+from threading import Thread
+from typing import Any, Awaitable, Iterable, Optional, Union, cast
+
+import RNS
+if find_spec("aiohttp") is None:
+    aiohttp_stub = types.ModuleType("aiohttp")
+
+    class ClientSession:  # pylint: disable=too-few-public-methods
+        """Fallback aiohttp ClientSession used for pytak import-time typing."""
+
+    aiohttp_stub.ClientSession = ClientSession
+    sys.modules.setdefault("aiohttp", aiohttp_stub)
+import pytak
+
+from . import Event
+
+CotPayload = Union[Event, ET.Element, str, bytes, dict]
+
+
+def _shutdown_weak(ref: "weakref.ReferenceType[PytakClient]") -> None:
+    """Invoke shutdown on a weakly referenced :class:`PytakClient`."""
+
+    client = ref()
+    if client is None:
+        return
+    client._shutdown_sync()  # pylint: disable=protected-access
+
+
+def _is_iterable_payload(obj: Any) -> bool:
+    """Return True when the object should be treated as a payload collection."""
+    if isinstance(obj, (Event, ET.Element, str, bytes, dict)):
+        return False
+    return isinstance(obj, Iterable)
+
+
+def _payload_to_xml_bytes(payload: CotPayload) -> bytes:
+    """Convert supported payload types into ATAK XML bytes."""
+    if isinstance(payload, Event):
+        return payload.to_xml_bytes()
+    if isinstance(payload, ET.Element):
+        return ET.tostring(payload, encoding="utf-8")
+    if isinstance(payload, bytes):
+        return payload
+    if isinstance(payload, str):
+        return payload.encode("utf-8")
+    if isinstance(payload, dict):
+        return Event.from_dict(payload).to_xml_bytes()
+    raise TypeError(f"Unsupported payload type: {type(payload)!r}")
+
+
+class SendWorker(pytak.QueueWorker):
+    """pyTAK worker that pushes ATAK CoT XML payloads onto the TX queue."""
+
+    def __init__(
+        self,
+        queue: asyncio.Queue,
+        config: SectionProxy,
+        message: Union[CotPayload, Iterable[CotPayload]],
+    ) -> None:
+        super().__init__(queue, config)
+        # Ensure a concrete list of CotPayload so the type checker knows
+        # iterating yields a CotPayload for handle_data(...)
+        self._messages: list[CotPayload]
+        if _is_iterable_payload(message):
+            self._messages = list(cast(Iterable[CotPayload], message))
+        else:
+            self._messages = [cast(CotPayload, message)]
+
+    async def handle_data(self, data: CotPayload) -> None:
+        await self.put_queue(_payload_to_xml_bytes(data))
+
+    async def run(self, number_of_iterations: int = 0):
+        _ = number_of_iterations
+        for payload in self._messages:
+            await self.handle_data(payload)
+
+
+class ReceiveWorker(pytak.QueueWorker):
+    """pyTAK worker that optionally parses incoming CoT XML into Event objects."""
+
+    def __init__(
+        self, queue: asyncio.Queue, config: SectionProxy, parse: bool = True
+    ) -> None:
+        super().__init__(queue, config)
+        self._parse = parse
+        # store parsed or raw data here so callers can inspect worker instances
+        self.result: Optional[Any] = None
+
+    async def handle_data(self, data: Any) -> None:
+        """Parse queue data into an Event when requested."""
+
+        if not self._parse:
+            self.result = data
+            return
+        try:
+            self.result = Event.from_xml(data)
+        except (ET.ParseError, TypeError, ValueError, AttributeError):
+            self.result = data
+
+    async def run(self, number_of_iterations: int = 0) -> None:
+        _ = number_of_iterations
+        try:
+            data = await self.queue.get()
+        except (asyncio.CancelledError, RuntimeError):
+            return None
+        await self.handle_data(data)
+        return None
+
+
+class StreamSendWorker(SendWorker):
+    """Continuous send worker that drains an outbound queue."""
+
+    def __init__(
+        self,
+        queue: asyncio.Queue,
+        config: SectionProxy,
+        outbound_queue: asyncio.Queue,
+        stop_event: asyncio.Event,
+    ) -> None:
+        super().__init__(queue, config, [])
+        self._outbound_queue = outbound_queue
+        self._stop_event = stop_event
+
+    async def run(self, number_of_iterations: int = 0):
+        iterations = 0
+        while not self._stop_event.is_set():
+            if number_of_iterations and iterations >= number_of_iterations:
+                return None
+            try:
+                payload = await asyncio.wait_for(
+                    self._outbound_queue.get(), timeout=0.2
+                )
+            except asyncio.TimeoutError:
+                continue
+            except (asyncio.CancelledError, RuntimeError):
+                return None
+            await self.handle_data(payload)
+            iterations += 1
+
+
+class FTSCLITool(pytak.CLITool):
+    """PyTAK CLI tool wrapper that tracks coroutine tasks for testing."""
+
+    def __init__(
+        self,
+        config: ConfigParser,
+        tx_queue: Union[asyncio.Queue, None] = None,
+        rx_queue: Union[asyncio.Queue, None] = None,
+    ) -> None:
+        self.config_parser = config if isinstance(config, ConfigParser) else None
+        section: ConfigParser | SectionProxy
+        if isinstance(config, ConfigParser):
+            section = (
+                config[config.sections()[0]] if config.sections() else config["DEFAULT"]
+            )
+        else:
+            section = config
+        super().__init__(section, tx_queue, rx_queue)
+        self.section = section
+        self.tasks_to_complete = set()
+        self.running_c_tasks = set()
+        # store results from the last run here
+        self.results: list[Any] = []
+
+    def add_c_task(self, task):
+        """Register a coroutine worker task to run alongside pyTAK tasks."""
+
+        self.tasks_to_complete.add(task)
+
+    def run_c_task(self, task):
+        """Schedule a coroutine worker task and keep a handle for teardown."""
+
+        self.running_c_tasks.add(asyncio.ensure_future(task.run()))
+
+    def run_c_tasks(self, tasks=None):
+        """Schedule all coroutine worker tasks."""
+
+        tasks = tasks or self.tasks_to_complete
+        for task in tasks:
+            self.run_c_task(task)
+
+    async def setup(self) -> None:
+        """Connect to the configured TAK server and log outcomes."""
+
+        cot_url = self.config.get("COT_URL", "")
+        try:
+            await super().setup()
+        except Exception as exc:  # pylint: disable=broad-exception-caught
+            self._logger.error(
+                "Failed to connect to TAK server at %s: %s", cot_url or "unknown", exc
+            )
+            RNS.log(
+                f"Failed to connect to TAK server at {cot_url or 'unknown'}: {exc}",
+                RNS.LOG_ERROR,
+            )
+            raise
+        self._logger.info("Connected to TAK server at %s", cot_url or "unknown")
+        RNS.log(f"Connected to TAK server at {cot_url or 'unknown'}", RNS.LOG_INFO)
+
+    async def run(self, number_of_iterations: int = 0) -> None:
+        """Runs this Thread and its associated coroutine tasks."""
+        _ = number_of_iterations
+        self._logger.info("Run: %s", self.__class__)
+
+        self.run_tasks()
+        self.run_c_tasks()
+
+        _done, _ = await asyncio.wait(
+            self.running_c_tasks, return_when=asyncio.ALL_COMPLETED
+        )
+
+        # Give the TX/RX workers a moment to drain the queues before cancelling
+        # them. Without this pause, the main loop could cancel the TX worker
+        # before it flushes the enqueued CoT payload.
+        await asyncio.sleep(getattr(self, "min_period", 0.1) or 0.1)
+
+        results: list[Any] = []
+
+        # Collect results from worker instances (ReceiveWorker stores parsed data
+        # on .result) instead of relying on coroutine return values.
+        for task in self.tasks_to_complete:
+            res = getattr(task, "result", None)
+            if res is not None:
+                results.append(res)
+
+        # Close TX and RX workers aka connection to Server
+        for task in self.running_tasks:
+            task.cancel()
+
+        # store results on the instance and return None to match base class
+        self.results = results
+        return None
+
+
+class PytakWorkerManager:  # pylint: disable=too-many-instance-attributes
+    """Manage a persistent PyTAK CLI tool and worker queue."""
+
+    def __init__(
+        self, cli_tool: FTSCLITool, section: SectionProxy, parse_inbound: bool
+    ) -> None:
+        self.cli_tool = cli_tool
+        self.section = section
+        self.parse_inbound = parse_inbound
+        self._outbound: asyncio.Queue = asyncio.Queue()
+        self._stop_event = asyncio.Event()
+        self._results: list[Any] = []
+        self._task: Optional[asyncio.Task] = None
+        self._session_task: Optional[asyncio.Task] = None
+        self._logger = getattr(cli_tool, "_logger", logging.getLogger(__name__))
+        self._backoff_seconds = 1.0
+
+    async def start(self) -> None:
+        """Start the long-running PyTAK session if it is not active."""
+
+        if self._stop_event.is_set():
+            self._stop_event = asyncio.Event()
+        if self._task is None or self._task.done():
+            self._task = asyncio.create_task(self._run_session())
+
+    async def stop(self) -> None:
+        """Stop the PyTAK session and cancel worker tasks."""
+
+        self._stop_event.set()
+        if self._task is not None:
+            self._task.cancel()
+            with suppress(asyncio.CancelledError):
+                await self._task
+            self._task = None
+            self._session_task = None
+
+    async def enqueue(self, message: CotPayload) -> None:
+        """Queue a payload for transmission over the active session."""
+
+        await self._outbound.put(message)
+
+    def results(self) -> list[Any]:
+        """Return results collected from the most recent receive worker."""
+
+        return list(self._results)
+
+    async def _run_session(self) -> None:
+        """Run a PyTAK session with exponential backoff on failures."""
+        while not self._stop_event.is_set():
+            send_stop = asyncio.Event()
+            try:
+                await self.cli_tool.setup()
+                self.cli_tool.tasks_to_complete.clear()
+                self.cli_tool.running_c_tasks.clear()
+
+                send_worker = StreamSendWorker(
+                    cast(asyncio.Queue, self.cli_tool.tx_queue),
+                    self.section,
+                    self._outbound,
+                    send_stop,
+                )
+                receive_worker = ReceiveWorker(
+                    cast(asyncio.Queue, self.cli_tool.rx_queue),
+                    self.section,
+                    parse=self.parse_inbound,
+                )
+
+                self.cli_tool.add_c_task(send_worker)
+                self.cli_tool.add_c_task(receive_worker)
+                self._results.clear()
+
+                self._session_task = asyncio.create_task(self.cli_tool.run())
+                try:
+                    await self._session_task
+                finally:
+                    send_stop.set()
+                    if self._session_task:
+                        self._session_task.cancel()
+                        with suppress(asyncio.CancelledError):
+                            await self._session_task
+
+                if getattr(receive_worker, "result", None) is not None:
+                    self._results.append(receive_worker.result)
+            except asyncio.CancelledError:
+                send_stop.set()
+                if self._session_task is not None:
+                    self._session_task.cancel()
+                    with suppress(asyncio.CancelledError):
+                        await self._session_task
+                raise
+            except Exception as exc:  # pragma: no cover - defensive logging  # pylint: disable=broad-exception-caught
+                send_stop.set()
+                self._logger.error("PyTAK session error: %s", exc)
+                await asyncio.sleep(self._backoff_seconds)
+                self._backoff_seconds = min(self._backoff_seconds * 2, 30.0)
+            else:
+                send_stop.set()
+                self._backoff_seconds = 1.0
+        return None
+
+
+class PytakClient:  # pylint: disable=too-many-instance-attributes
+    """Utility wrapper that wires ATAK Event payloads into pyTAK workers."""
+
+    def __init__(self, config: Optional[ConfigParser] = None) -> None:
+        self._config = config
+        self._cli_tool: Optional[FTSCLITool] = None
+        self._worker_manager: Optional[PytakWorkerManager] = None
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._loop_thread: Optional[Thread] = None
+        self._loop_ready = ThreadEvent()
+        self._loop_lock = Lock()
+        atexit.register(_shutdown_weak, weakref.ref(self))
+
+    def __del__(self) -> None:
+        try:
+            self._shutdown_sync()  # pylint: disable=protected-access
+        except Exception:  # pylint: disable=broad-exception-caught
+            pass
+
+    def _setup_config(self) -> ConfigParser:
+        """Create config if a custom one is not passed."""
+        config = ConfigParser()
+        config["fts"] = {
+            "COT_URL": "tcp://127.0.0.1:8087",
+            "CALLSIGN": "FTS_PYTAK",
+            "TAK_PROTO": "0",
+            "FTS_COMPAT": "1",
+        }
+        return config
+
+    def _ensure_config(self, config: Optional[ConfigParser]) -> ConfigParser:
+        """
+        Ensure a configuration object is present for PyTAK workers.
+
+        Args:
+            config (ConfigParser | None): Custom configuration provided by the caller.
+
+        Returns:
+            ConfigParser: The configuration to use for PyTAK interactions.
+        """
+        if config is not None:
+            if self._config is None:
+                self._config = config
+            return config
+        if self._config is None:
+            self._config = self._setup_config()
+        return self._config
+
+    def _config_section(
+        self, config: ConfigParser, section: str = "fts"
+    ) -> SectionProxy:
+        """
+        Return the requested section or a fallback from a configuration object.
+
+        Args:
+            config (ConfigParser): Configuration containing PyTAK settings.
+            section (str): Desired section name. Defaults to ``"fts"``.
+
+        Returns:
+            SectionProxy: Section with connection parameters.
+
+        Raises:
+            ValueError: If the configuration has no sections.
+        """
+        if config.has_section(section):
+            return config[section]
+        sections = config.sections()
+        if sections:
+            return config[sections[0]]
+        raise ValueError("Configuration must contain at least one section.")
+
+    def _ensure_cli_tool(self, config: ConfigParser) -> FTSCLITool:
+        """Create or return a cached CLI tool backed by shared queues."""
+
+        if self._cli_tool is None:
+            tx_queue: asyncio.Queue = asyncio.Queue()
+            rx_queue: asyncio.Queue = asyncio.Queue()
+            self._cli_tool = FTSCLITool(config, tx_queue, rx_queue)
+        return self._cli_tool
+
+    def _ensure_manager(
+        self, config: ConfigParser, parse_inbound: bool
+    ) -> "PytakWorkerManager":
+        """
+        Return a running worker manager with the provided configuration.
+
+        Args:
+            config (ConfigParser): PyTAK configuration to apply.
+            parse_inbound (bool): Whether inbound CoT data should be parsed.
+
+        Returns:
+            PytakWorkerManager: The configured worker manager.
+        """
+
+        cli_tool = self._ensure_cli_tool(config)
+        if self._worker_manager is None:
+            section = self._config_section(config)
+            self._worker_manager = PytakWorkerManager(cli_tool, section, parse_inbound)
+        else:
+            self._worker_manager.parse_inbound = parse_inbound
+        return self._worker_manager
+
+    async def create_and_send_message(
+        self,
+        message: Union[CotPayload, Iterable[CotPayload]],
+        config: Optional[ConfigParser] = None,
+        parse_inbound: bool = True,
+    ) -> list[Any]:
+        """
+        Send one or more CoT payloads through a PyTAK worker session.
+
+        Args:
+            message (CotPayload | Iterable[CotPayload]): Payload(s) to dispatch.
+            config (ConfigParser | None): Optional configuration override.
+            parse_inbound (bool): Whether to parse inbound data into :class:`Event`.
+
+        Returns:
+            list[Any]: Parsed or raw results from the receive worker.
+        """
+        cfg = self._ensure_config(config)
+        manager = self._ensure_manager(cfg, parse_inbound)
+        await self._run_in_loop(manager.start())
+        await self._run_in_loop(manager.enqueue(message))
+        return manager.results()
+
+    async def send_event(
+        self,
+        event: Event,
+        config: Optional[ConfigParser] = None,
+        parse_inbound: bool = True,
+    ):
+        """Convenience helper that sends a single Event."""
+        return await self.create_and_send_message(
+            event, config=config, parse_inbound=parse_inbound
+        )
+
+    @staticmethod
+    def _start_loop(
+        loop: asyncio.AbstractEventLoop, ready_event: ThreadEvent
+    ) -> None:
+        """
+        Start the event loop on a dedicated thread and signal readiness.
+
+        Args:
+            loop (asyncio.AbstractEventLoop): Event loop to run.
+            ready_event (ThreadEvent): Event set once the loop is running.
+        """
+        asyncio.set_event_loop(loop)
+        ready_event.set()
+        loop.run_forever()
+
+    def _ensure_loop(self) -> asyncio.AbstractEventLoop:
+        """Ensure a background event loop exists for PyTAK tasks."""
+
+        with self._loop_lock:
+            if self._loop is not None and self._loop.is_running():
+                return self._loop
+            loop = asyncio.new_event_loop()
+            self._loop = loop
+            self._loop_ready.clear()
+            thread = Thread(
+                target=self._start_loop, args=(loop, self._loop_ready), daemon=True
+            )
+            self._loop_thread = thread
+            thread.start()
+        self._loop_ready.wait()
+        return cast(asyncio.AbstractEventLoop, self._loop)
+
+    async def _run_in_loop(self, coro: Awaitable[Any]) -> Any:
+        """Execute a coroutine on the dedicated event loop and await it."""
+
+        loop = self._ensure_loop()
+        try:
+            running_loop = asyncio.get_running_loop()
+        except RuntimeError:
+            running_loop = None
+        if running_loop is loop:
+            return await coro
+        future = asyncio.run_coroutine_threadsafe(coro, loop)
+        return await asyncio.wrap_future(future)
+
+    async def stop(self) -> None:
+        """Stop the PyTAK worker manager and background loop."""
+
+        if self._worker_manager is not None:
+            await self._run_in_loop(self._worker_manager.stop())
+            self._worker_manager = None
+        if self._loop is not None and self._loop.is_running():
+            self._loop.call_soon_threadsafe(self._loop.stop)
+            if self._loop_thread is not None:
+                self._loop_thread.join(timeout=1.0)
+        self._loop = None
+        self._loop_thread = None
+
+    def _shutdown_sync(self) -> None:
+        """Best-effort cleanup for interpreter shutdown or GC."""
+
+        if self._loop is None or not self._loop.is_running():
+            self._loop = None
+            self._loop_thread = None
+            self._worker_manager = None
+            return
+
+        if self._worker_manager is not None:
+            try:
+                future = asyncio.run_coroutine_threadsafe(
+                    self._worker_manager.stop(), self._loop
+                )
+                future.result(timeout=1.0)
+            except Exception:  # pylint: disable=broad-exception-caught
+                pass
+            self._worker_manager = None
+
+        self._loop.call_soon_threadsafe(self._loop.stop)
+        if self._loop_thread is not None:
+            self._loop_thread.join(timeout=1.0)
+        self._loop = None
+        self._loop_thread = None