From 14a2d0ae12e386262eba1086ed836a7eebaff541 Mon Sep 17 00:00:00 2001
From: InsanePrawn <insane.prawny@gmail.com>
Date: Thu, 17 Feb 2022 16:25:31 +0100
Subject: [PATCH] wrapper/: introduce Wrapper protocol

---
 wrapper/__init__.py | 15 +++++++++------
 wrapper/docker.py   |  2 +-
 wrapper/wrapper.py  | 32 +++++++++++++++++++++++++++-----
 3 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/wrapper/__init__.py b/wrapper/__init__.py
index 48a3f2e..9845008 100644
--- a/wrapper/__init__.py
+++ b/wrapper/__init__.py
@@ -1,13 +1,13 @@
-import os
 import click
 import logging
 
 from config import config
 from utils import programs_available
 from .docker import DockerWrapper
+from .wrapper import Wrapper
 
-wrapper_impls = {
-    'docker': DockerWrapper,
+wrapper_impls: dict[str, Wrapper] = {
+    'docker': DockerWrapper(),
 }
 
 
@@ -15,15 +15,18 @@ def get_wrapper_type(wrapper_type: str = None):
     return wrapper_type or config.file['wrapper']['type']
 
 
+def get_wrapper_impl(wrapper_type: str = None) -> Wrapper:
+    return wrapper_impls[get_wrapper_type(wrapper_type)]
+
+
 def wrap(wrapper_type: str = None):
     wrapper_type = get_wrapper_type(wrapper_type)
     if wrapper_type != 'none':
-        wrapper_impls[wrapper_type]().wrap()
+        get_wrapper_impl(wrapper_type).wrap()
 
 
 def is_wrapped(wrapper_type: str = None):
-    wrapper_type = get_wrapper_type(wrapper_type)
-    return os.getenv('KUPFERBOOTSTRAP_WRAPPED') == wrapper_type.capitalize()
+    return get_wrapper_impl(wrapper_type).is_wrapped()
 
 
 def enforce_wrap(no_wrapper=False):
diff --git a/wrapper/docker.py b/wrapper/docker.py
index 8119856..7ca06c3 100644
--- a/wrapper/docker.py
+++ b/wrapper/docker.py
@@ -19,7 +19,7 @@ def docker_volumes_args(volume_mappings: dict[str, str]) -> list[str]:
 
 
 class DockerWrapper(BaseWrapper):
-    type = 'docker'
+    type: str = 'docker'
 
     def wrap(self):
         script_path = config.runtime['script_source_dir']
diff --git a/wrapper/wrapper.py b/wrapper/wrapper.py
index 1abb38d..37f0787 100644
--- a/wrapper/wrapper.py
+++ b/wrapper/wrapper.py
@@ -3,15 +3,33 @@ import os
 import uuid
 import pathlib
 
+from typing import Protocol
+
 from config import config, dump_file as dump_config_file
 from constants import CHROOT_PATHS
 
 
-class BaseWrapper:
-    id = None
-    identifier = None
-    type = None
-    wrapped_config_path = None
+class Wrapper(Protocol):
+    """Wrappers wrap kupferbootstrap in some form of isolation from the host OS, i.e. docker or chroots"""
+
+    def wrap(self):
+        """Instructs the wrapper to reexecute kupferbootstrap in a wrapped environment"""
+
+    def stop(self):
+        """Instructs the wrapper to stop the wrapped instance and clean up"""
+
+    def is_wrapped(self) -> bool:
+        """
+        Queries the wrapper whether it believes we're executing wrapped by it currently.
+        Checks `env[KUPFERBOOTSTRAP_WRAPPED] == self.type.capitalize()` by default.
+        """
+
+
+class BaseWrapper(Wrapper):
+    id: str
+    identifier: str
+    type: str
+    wrapped_config_path: str
 
     def __init__(self, random_id: str = None, name: str = None):
         self.uuid = str(random_id or uuid.uuid4())
@@ -49,6 +67,7 @@ class BaseWrapper:
     ) -> str:
         wrapped_config = f'{target_path.rstrip("/")}/{self.identifier}_wrapped.toml'
 
+        # FIXME: these at_exit hooks should go and be called from somewhere better suited
         def at_exit():
             self.stop()
             os.remove(wrapped_config)
@@ -73,6 +92,9 @@ class BaseWrapper:
     def stop(self):
         raise NotImplementedError()
 
+    def is_wrapped(self):
+        return os.getenv('KUPFERBOOTSTRAP_WRAPPED') == self.type.capitalize()
+
     def get_bind_mounts_default(self, wrapped_config_path: str = None, ssh_dir: str = None, target_home: str = '/root'):
         wrapped_config_path = wrapped_config_path or self.wrapped_config_path
         ssh_dir = ssh_dir or os.path.join(pathlib.Path.home(), '.ssh')