From 5cba20c20a210e62c2133f3290be551a07d66e94 Mon Sep 17 00:00:00 2001
From: Chris Peterson <chris.peterson@canonical.com>
Date: Sat, 2 Nov 2024 18:24:13 +0100
Subject: [PATCH] utils: introduce passlib

Transient commit to show passlib functionality is equivalent with
crypt module, given the right parameters. Rebase me before merge.
---
 snapcraft.yaml                    |  1 +
 subiquitycore/tests/test_utils.py | 32 +++++++++++++++++++++++++++++++
 subiquitycore/utils.py            | 24 +++++++++++++++++++++++
 3 files changed, 57 insertions(+)

diff --git a/snapcraft.yaml b/snapcraft.yaml
index 07e275a59..a077d524e 100644
--- a/snapcraft.yaml
+++ b/snapcraft.yaml
@@ -136,6 +136,7 @@ parts:
       - python3-minimal
       - python3-more-itertools
       - python3-oauthlib
+      - python3-passlib
       - python3-pkg-resources
       - python3-pyroute2
       - python3-pyrsistent
diff --git a/subiquitycore/tests/test_utils.py b/subiquitycore/tests/test_utils.py
index b2809f780..6f7207d73 100644
--- a/subiquitycore/tests/test_utils.py
+++ b/subiquitycore/tests/test_utils.py
@@ -20,8 +20,10 @@
 from subiquitycore.tests import SubiTestCase
 from subiquitycore.utils import (
     _zsys_uuid_charset,
+    crypt_password,
     gen_zsys_uuid,
     orig_environ,
+    passlib_crypt,
     system_scripts_env,
 )
 
@@ -129,3 +131,33 @@ def test_zsys_uuid(self):
         for i in range(10):
             uuid = gen_zsys_uuid()
             self.assertEqual(6, len(uuid), uuid)
+
+
+class TestCryptPassword(SubiTestCase):
+    @patch("subiquitycore.utils._generate_salt")
+    def test_compare_passlib_with_crypt(self, salt_mock):
+        """Test passlib module output is equivalent with python crypt module."""
+
+        # Test SHA-512
+        salt_mock.return_value = "mock.salt"
+        python = crypt_password("ubuntu", "SHA-512")
+        passlib = passlib_crypt("ubuntu", "SHA-512")
+        self.assertEqual(python, passlib)
+
+        # Test SHA-256
+        salt_mock.return_value = "mock.salt"
+        python = crypt_password("ubuntu", "SHA-256")
+        passlib = passlib_crypt("ubuntu", "SHA-256")
+        self.assertEqual(python, passlib)
+
+        # Test MD5
+        salt_mock.return_value = "mock.salt"
+        python = crypt_password("ubuntu", "MD5")
+        passlib = passlib_crypt("ubuntu", "MD5")
+        self.assertEqual(python, passlib)
+
+        # Test DES
+        salt_mock.return_value = "mock.salt"
+        python = crypt_password("ubuntu", "DES")
+        passlib = passlib_crypt("ubuntu", "DES")
+        self.assertEqual(python, passlib)
diff --git a/subiquitycore/utils.py b/subiquitycore/utils.py
index 26534f996..058ba3358 100644
--- a/subiquitycore/utils.py
+++ b/subiquitycore/utils.py
@@ -21,6 +21,8 @@
 import tempfile
 from typing import Any, Dict, List, Sequence
 
+import passlib.hash
+
 log = logging.getLogger("subiquitycore.utils")
 
 
@@ -257,6 +259,28 @@ def crypt_password(passwd, algo="SHA-512"):
     return crypt.crypt(passwd, algos[algo] + salt)
 
 
+def passlib_crypt(passwd, algo="SHA-512"):
+    # Use rounds=5000 where possible to be equivalent w/ crypt.
+    algos = {
+        "SHA-512": passlib.hash.sha512_crypt.using(rounds=5000),
+        "SHA-256": passlib.hash.sha256_crypt.using(rounds=5000),
+        "MD5": passlib.hash.md5_crypt,
+        "DES": passlib.hash.des_crypt,
+    }
+
+    salt = _generate_salt()
+
+    # MD5 only supports salts of <= 8 characters
+    if algo == "MD5":
+        salt = salt[:8]
+    # DES only supports salts of <= 2 charactes
+    elif algo == "DES":
+        salt = salt[:2]
+
+    handler = algos[algo].using(salt=salt)
+    return handler.hash(passwd)
+
+
 def disable_subiquity():
     """Stop subiquity service; which also restores getty service"""
     log.info("disabling subiquity service")