From 2c25f6601510c6157bf1575d38541d4395989548 Mon Sep 17 00:00:00 2001
From: Andrew Walker <awalker@ixsystems.com>
Date: Tue, 27 Aug 2024 13:50:59 -0600
Subject: [PATCH] Add basic utilities for manipulating pam_tdb data

---
 src/middlewared/middlewared/utils/crypto.py   |  15 +++
 src/middlewared/middlewared/utils/tdb.py      |   6 +
 .../middlewared/utils/user_api_key.py         | 112 ++++++++++++++++++
 3 files changed, 133 insertions(+)
 create mode 100644 src/middlewared/middlewared/utils/user_api_key.py

diff --git a/src/middlewared/middlewared/utils/crypto.py b/src/middlewared/middlewared/utils/crypto.py
index 25880fad0ce34..5f5d826f29991 100644
--- a/src/middlewared/middlewared/utils/crypto.py
+++ b/src/middlewared/middlewared/utils/crypto.py
@@ -1,3 +1,5 @@
+from base64 import b64encode
+from hashlib import pbkdf2_hmac
 from secrets import choice, compare_digest, token_urlsafe, token_hex
 from string import ascii_letters, digits, punctuation
 
@@ -63,3 +65,16 @@ def generate_nt_hash(passwd):
     """
     md4_hash_bytes = md4_hash_blob(passwd.encode('utf-16le'))
     return md4_hash_bytes.hex().upper()
+
+
+def generate_pbkdf2_512(passwd):
+    """
+    Generate a pbkdf2_sha512 hash for password. This is used for
+    verification of API keys.
+    """
+    prefix = 'pbkdf2-sha512'
+    rounds = 500000
+    salt_length = 16
+    salt = generate_string(string_size=salt_length, extra_chars='./').encode()
+    hash = pbkdf2_hmac('sha512', passwd.encode(), salt, rounds)
+    return f'${prefix}${rounds}${b64encode(salt).decode()}${b64encode(hash).decode()}'
diff --git a/src/middlewared/middlewared/utils/tdb.py b/src/middlewared/middlewared/utils/tdb.py
index d67ef99317006..974990d8b6169 100644
--- a/src/middlewared/middlewared/utils/tdb.py
+++ b/src/middlewared/middlewared/utils/tdb.py
@@ -100,6 +100,12 @@ class TDBHandle:
     opath_fd = FD_CLOSED
     keys_null_terminated = False
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, tp, val, traceback):
+        self.close()
+
     def close(self):
         """ Close the TDB handle and O_PATH open for the file """
         if self.opath_fd == FD_CLOSED and self.hdl is None:
diff --git a/src/middlewared/middlewared/utils/user_api_key.py b/src/middlewared/middlewared/utils/user_api_key.py
new file mode 100644
index 0000000000000..673a4c7a8c1ee
--- /dev/null
+++ b/src/middlewared/middlewared/utils/user_api_key.py
@@ -0,0 +1,112 @@
+import os
+
+from base64 import b64encode
+from dataclasses import dataclass
+from struct import pack
+from uuid import uuid4
+from .tdb import (
+    TDBDataType,
+    TDBHandle,
+    TDBOptions,
+    TDBPathType,
+)
+
+
+PAM_TDB_DIR = '/var/run/pam_tdb'
+PAM_TDB_FILE = os.path.join(PAM_TDB_DIR, 'pam_tdb.tdb')
+PAM_TDB_DIR_MODE = 0o700
+PAM_TDB_VERSION = 1
+PAM_TDB_MAX_KEYS = 10  # Max number of keys per user
+
+PAM_TDB_OPTIONS = TDBOptions(TDBPathType.CUSTOM, TDBDataType.BYTES)
+
+
+@dataclass(frozen=True)
+class UserApiKey:
+    expiry: int
+    userhash: str
+
+
+@dataclass(frozen=True)
+class PamTdbEntry:
+    keys: list[UserApiKey]
+    username: str
+
+
+def _setup_pam_tdb_dir() -> None:
+    os.makedirs(PAM_TDB_DIR, mode=PAM_TDB_DIR_MODE, exist_ok=True)
+    os.chmod(PAM_TDB_DIR, PAM_TDB_DIR_MODE)
+
+
+def _pack_user_auth_key(api_key: UserApiKey) -> bytes:
+    """
+    Convert UserAuthToken to bytes for TDB insertion.
+    This is packed struct with expiry converted into signed 64 bit
+    integer, and the userhash converted into a pascal string.
+    """
+    if not isinstance(api_key, UserApiKey):
+        raise TypeError(f'{type(api_key)}: not a UserAuthToken')
+
+    userhash = api_key.userhash.encode() + b'\x00'
+    return pack(f'<q{len(userhash)}p', api_key.expiry, userhash)
+
+
+def write_entry(hdl: TDBHandle, entry: PamTdbEntry) -> None:
+    """
+    Convert PamTdbEntry object into a packed struct and insert
+    into tdb file.
+
+    key: username
+    value: uint32_t (version) + uint32_t (cnt of keys)
+    """
+    if not isinstance(entry, PamTdbEntry):
+        raise TypeError(f'{type(entry)}: expected UserAuthToken')
+
+    key_cnt = len(entry.keys)
+    if key_cnt > PAM_TDB_MAX_KEYS:
+        raise ValueError(f'{key_cnt}: count of entries exceeds maximum')
+
+    entry_bytes = pack('<II', PAM_TDB_VERSION, len(entry.keys))
+    parsed_cnt = 0
+    for key in entry.keys:
+        entry_bytes += _pack_user_auth_key(key)
+        parsed_cnt += 1
+
+    # since we've already packed struct with array length
+    # we need to rigidly ensure we don't exceed it.
+    assert parsed_cnt == key_cnt
+    hdl.store(entry.username, b64encode(entry_bytes))
+
+
+def flush_user_api_keys(pam_entries: list[PamTdbEntry]) -> None:
+    """
+    Write a PamTdbEntry object to the pam_tdb file for user
+    authentication. This method first writes to temporary file
+    and then renames over pam_tdb file to ensure flush is atomic
+    and reduce risk of lock contention while under a transaction
+    lock.
+
+    raises:
+        TypeError - not PamTdbEntry
+        AssertionError - count of entries changed while generating
+            tdb payload
+        RuntimeError - TDB library error
+    """
+    _setup_pam_tdb_dir()
+
+    if not isinstance(pam_entries, list):
+        raise TypeError('Expected list of PamTdbEntry objects')
+
+    tmp_path = os.path.join(PAM_TDB_DIR, f'tmp_{uuid4()}.tdb')
+
+    with TDBHandle(tmp_path, PAM_TDB_OPTIONS) as hdl:
+        hdl.keys_null_terminated = False
+
+        try:
+            for entry in pam_entries:
+                write_entry(hdl, entry)
+        except Exception:
+            os.remove(tmp_path)
+            raise
+
+    os.rename(tmp_path, PAM_TDB_FILE)