TOTP implementation

5 files changed, 141 insertions, 1 deletions

diff --git a/kvmd/apps/__init__.py b/kvmd/apps/__init__.py
index 7d786bc4..cbc2400e 100644
--- a/kvmd/apps/__init__.py
+++ b/kvmd/apps/__init__.py
@@ -365,6 +365,12 @@ def _get_config_scheme() -> dict:
                     "type": Option("", type=valid_stripped_string),
                     # Dynamic content
                 },
+
+                "totp": {
+                    "secret": {
+                        "file": Option("/etc/kvmd/totp.secret", type=valid_abs_path, if_empty=""),
+                    },
+                },
             },
 
             "info": {  # Accessed via global config, see kvmd/info for details
diff --git a/kvmd/apps/kvmd/__init__.py b/kvmd/apps/kvmd/__init__.py
index 4d891f40..a90cd30b 100644
--- a/kvmd/apps/kvmd/__init__.py
+++ b/kvmd/apps/kvmd/__init__.py
@@ -82,6 +82,8 @@ def main(argv: (list[str] | None)=None) -> None:
 
             external_type=config.auth.external.type,
             external_kwargs=(config.auth.external._unpack(ignore=["type"]) if config.auth.external.type else {}),
+
+            totp_secret_path=config.auth.totp.secret.file,
         ),
         info_manager=InfoManager(global_config),
         log_reader=(LogReader() if config.log_reader.enabled else None),
diff --git a/kvmd/apps/kvmd/auth.py b/kvmd/apps/kvmd/auth.py
index 6ae696ad..02639eb0 100644
--- a/kvmd/apps/kvmd/auth.py
+++ b/kvmd/apps/kvmd/auth.py
@@ -21,6 +21,7 @@
 
 
 import secrets
+import pyotp
 
 from ...logging import get_logger
 
@@ -42,6 +43,8 @@ class AuthManager:
 
         external_type: str,
         external_kwargs: dict,
+
+        totp_secret_path: str,
     ) -> None:
 
         self.__enabled = enabled
@@ -53,12 +56,14 @@ class AuthManager:
             self.__internal_service = get_auth_service_class(internal_type)(**internal_kwargs)
             get_logger().info("Using internal auth service %r", self.__internal_service.get_plugin_name())
 
+        self.__force_internal_users = force_internal_users
+
         self.__external_service: (BaseAuthService | None) = None
         if enabled and external_type:
             self.__external_service = get_auth_service_class(external_type)(**external_kwargs)
             get_logger().info("Using external auth service %r", self.__external_service.get_plugin_name())
 
-        self.__force_internal_users = force_internal_users
+        self.__totp_secret_path = totp_secret_path
 
         self.__tokens: dict[str, str] = {}  # {token: user}
 
@@ -71,6 +76,16 @@ class AuthManager:
         assert self.__enabled
         assert self.__internal_service
 
+        if self.__totp_secret_path:
+            with open(self.__totp_secret_path) as secret_file:
+                secret = secret_file.read().strip()
+            if secret:
+                code = passwd[-6:]
+                if not pyotp.TOTP(secret).verify(code):
+                    get_logger().error("Got access denied for user %r by TOTP", user)
+                    return False
+                passwd = passwd[:-6]
+
         if user not in self.__force_internal_users and self.__external_service:
             service = self.__external_service
         else:
diff --git a/kvmd/apps/totp/__init__.py b/kvmd/apps/totp/__init__.py
new file mode 100644
index 00000000..64be9b90
--- /dev/null
+++ b/kvmd/apps/totp/__init__.py
@@ -0,0 +1,93 @@
+# ========================================================================== #
+#                                                                            #
+#    KVMD - The main PiKVM daemon.                                           #
+#                                                                            #
+#    Copyright (C) 2018-2022  Maxim Devaev <[email protected]>               #
+#                                                                            #
+#    This program is free software: you can redistribute it and/or modify    #
+#    it under the terms of the GNU General Public License as published by    #
+#    the Free Software Foundation, either version 3 of the License, or       #
+#    (at your option) any later version.                                     #
+#                                                                            #
+#    This program is distributed in the hope that it will be useful,         #
+#    but WITHOUT ANY WARRANTY; without even the implied warranty of          #
+#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the           #
+#    GNU General Public License for more details.                            #
+#                                                                            #
+#    You should have received a copy of the GNU General Public License       #
+#    along with this program.  If not, see <https://www.gnu.org/licenses/>.  #
+#                                                                            #
+# ========================================================================== #
+
+
+import argparse
+
+import pyotp
+import qrcode
+
+from ...yamlconf import Section
+
+from .. import init
+
+
+# =====
+def _get_secret_path(config: Section) -> str:
+    path: str = config.kvmd.auth.totp.secret.file
+    if len(path) == 0:
+        raise SystemExit("Error: TOTP file path is empty (i.e. it was disabled)")
+    return path
+
+
+def _read_secret(config: Section) -> str:
+    with open(_get_secret_path(config)) as file:
+        return file.read().strip()
+
+
+# =====
+def _cmd_init(config: Section, options: argparse.Namespace) -> None:
+    if not options.force:
+        if _read_secret(config):
+            raise SystemExit("Error: the TOTP secret already exists")
+    with open(_get_secret_path(config), "w") as file:
+        file.write(pyotp.random_base32())
+    _cmd_show(config, options)
+
+
+def _cmd_show(config: Section, _: argparse.Namespace) -> None:
+    secret = _read_secret(config)
+    if len(secret) == 0:
+        raise SystemExit("Error: TOTP secret is not configured")
+    uri = pyotp.totp.TOTP(secret).provisioning_uri(issuer_name="PiKVM")
+    qr = qrcode.QRCode()
+    qr.add_data(uri)
+    print()
+    print(uri)
+    print()
+    qr.print_ascii(invert=True)
+    print()
+
+
+# =====
+def main(argv: (list[str] | None)=None) -> None:
+    (parent_parser, argv, config) = init(
+        add_help=False,
+        cli_logging=True,
+        argv=argv,
+    )
+    parser = argparse.ArgumentParser(
+        prog="kvmd-totp",
+        description="Manage KVMD TOTP secret",
+        parents=[parent_parser],
+    )
+    parser.set_defaults(cmd=(lambda *_: parser.print_help()))
+    subparsers = parser.add_subparsers()
+
+    cmd_setup_parser = subparsers.add_parser("init", help="Generate and show TOTP secret with QR code")
+    cmd_setup_parser.add_argument("-f", "--force", action="store_true", help="Overwrite an existing secret")
+    cmd_setup_parser.set_defaults(cmd=_cmd_init)
+
+    cmd_show_parser = subparsers.add_parser("show", help="Show the current TOTP secret with QR code")
+    cmd_show_parser.set_defaults(cmd=_cmd_show)
+
+    options = parser.parse_args(argv[1:])
+    options.cmd(config, options)
diff --git a/kvmd/apps/totp/__main__.py b/kvmd/apps/totp/__main__.py
new file mode 100644
index 00000000..3849d1b9
--- /dev/null
+++ b/kvmd/apps/totp/__main__.py
@@ -0,0 +1,24 @@
+# ========================================================================== #
+#                                                                            #
+#    KVMD - The main PiKVM daemon.                                           #
+#                                                                            #
+#    Copyright (C) 2018-2022  Maxim Devaev <[email protected]>               #
+#                                                                            #
+#    This program is free software: you can redistribute it and/or modify    #
+#    it under the terms of the GNU General Public License as published by    #
+#    the Free Software Foundation, either version 3 of the License, or       #
+#    (at your option) any later version.                                     #
+#                                                                            #
+#    This program is distributed in the hope that it will be useful,         #
+#    but WITHOUT ANY WARRANTY; without even the implied warranty of          #
+#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the           #
+#    GNU General Public License for more details.                            #
+#                                                                            #
+#    You should have received a copy of the GNU General Public License       #
+#    along with this program.  If not, see <https://www.gnu.org/licenses/>.  #
+#                                                                            #
+# ========================================================================== #
+
+
+from . import main
+main()