1 files changed, 124 insertions, 0 deletions
diff --git a/desmo_api/fsm.py b/desmo_api/fsm.py
new file mode 100644
index 0000000..dfbef20
--- /dev/null
+++ b/desmo_api/fsm.py
@@ -0,0 +1,124 @@
+import asyncio
+from collections.abc import Coroutine
+import logging
+from statemachine import StateMachine, State
+
+
+from . import actions
+from . import hcloud_dns
+from . import db
+
+
+logger = logging.getLogger(__name__)
+
+
+class JailStateMachine(StateMachine):
+    def __init__(
+        self,
+        dns_client: hcloud_dns.HCloudDNS,
+        database: db.DB,
+        name: str,
+    ):
+        self._dns_client = dns_client
+        self._name = name
+        self._db = database
+        self._tasks = set()
+
+        super().__init__()
+
+    uninitialized = State(initial=True)
+    jail_provisioning = State()
+    dns_provisioning = State()
+    jail_setup = State()
+    jail_ready = State()
+    jail_removal = State()
+    dns_deprovisioning = State()
+    terminated = State()
+
+    initialize = uninitialized.to(jail_provisioning) | terminated.to(jail_provisioning)
+    jail_provisioned = jail_provisioning.to(dns_provisioning)
+    jail_provisioning_failed = jail_provisioning.to(jail_provisioning)
+    dns_provisioned = dns_provisioning.to(jail_setup)
+    dns_provisioning_failed = dns_provisioning.to(dns_provisioning)
+    jail_setup_done = jail_setup.to(jail_ready)
+    jail_setup_failed = jail_setup.to(jail_setup)
+    remove_jail = jail_ready.to(jail_removal)
+    jail_removal_failed = jail_removal.to(jail_removal)
+    jail_removed = jail_removal.to(dns_deprovisioning)
+    dns_deprovisioned = dns_deprovisioning.to(terminated)
+    dns_deprovisioning_failed = dns_deprovisioning.to(dns_deprovisioning)
+
+    def start_on_enter_task(self):
+        state = self.current_state.id
+        switch = {
+            "jail_provisioning": self.on_enter_jail_provisioning,
+            "dns_provisioning": self.on_enter_dns_provisioning,
+            "jail_setup": self.on_enter_jail_setup,
+            "jail_ready": self.on_enter_jail_ready,
+            "jail_removal": self.on_enter_jail_removal,
+            "dns_deprovisioning": self.on_enter_dns_deprovisioning,
+        }
+        func = switch.get(state)
+        if func is not None:
+            func()
+
+    async def _store_state_and_run(self, awaitable: Coroutine):
+        await self._db.set_jail_state(self._name, self.current_state.id)
+        await awaitable
+
+    def on_enter_jail_provisioning(self):
+        task = asyncio.create_task(
+            self._store_state_and_run(
+                actions.start_jail_provisioning(self, self._db, self._name)
+            )
+        )
+        self._tasks.add(task)
+
+    def on_enter_dns_provisioning(self):
+        task = asyncio.create_task(
+            self._store_state_and_run(
+                actions.start_dns_provisioning(
+                    self, self._db, self._dns_client, self._name
+                )
+            )
+        )
+        self._tasks.add(task)
+
+    def on_enter_jail_setup(self):
+        task = asyncio.create_task(
+            self._store_state_and_run(
+                actions.start_jail_setup(self, self._db, self._name)
+            )
+        )
+        self._tasks.add(task)
+
+    def on_enter_jail_ready(self):
+        task = asyncio.create_task(
+            self._store_state_and_run(
+                actions.start_jail_watch(self, self._db, self._name)
+            )
+        )
+        self._tasks.add(task)
+
+    def on_enter_jail_removal(self):
+        task = asyncio.create_task(
+            self._store_state_and_run(
+                actions.start_jail_removal(self, self._db, self._name)
+            )
+        )
+        self._tasks.add(task)
+
+    def on_enter_dns_deprovisioning(self):
+        task = asyncio.create_task(
+            self._store_state_and_run(
+                actions.start_dns_deprovisioning(
+                    self, self._dns_client, self._db, self._name
+                )
+            )
+        )
+        self._tasks.add(task)
+
+    def on_enter_terminated(self):
+        logger.info("Jail %s terminated", self._name)
+        for task in self._tasks:
+            task.cancel()