From feeae86f7885a8b2aeb92a24f1a22287a0c27157 Mon Sep 17 00:00:00 2001
From: Dmitry Volodin <dv@nocproject.org>
Date: Mon, 22 Jan 2024 17:33:16 +0100
Subject: [PATCH] BaseAuthKey.get_localized_key() helper function

---
 src/gufo/snmp/_fast.pyi |  3 +++
 src/gufo/snmp/user.py   | 18 +++++++++++++++++-
 src/lib.rs              |  1 +
 src/util.rs             | 19 +++++++++++++++++++
 tests/test_user.py      | 24 ++++++++++++++++++++++++
 5 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/src/gufo/snmp/_fast.pyi b/src/gufo/snmp/_fast.pyi
index d09a38c..fa939e5 100755
--- a/src/gufo/snmp/_fast.pyi
+++ b/src/gufo/snmp/_fast.pyi
@@ -122,3 +122,6 @@ class SnmpV3ClientSocket(object):
     def recv_refresh(self: "SnmpV3ClientSocket") -> None: ...
 
 def get_master_key(auth_alg: int, passwd: bytes) -> bytes: ...
+def get_localized_key(
+    auth_alg: int, passwd: bytes, engine_id: bytes
+) -> bytes: ...
diff --git a/src/gufo/snmp/user.py b/src/gufo/snmp/user.py
index a464131..2b590d9 100644
--- a/src/gufo/snmp/user.py
+++ b/src/gufo/snmp/user.py
@@ -11,7 +11,7 @@
 from typing import List, Optional, Type, TypeVar
 
 # Gufo SNMP modules
-from ._fast import get_master_key
+from ._fast import get_localized_key, get_master_key
 
 K = TypeVar("K", bound="BaseKey")
 
@@ -103,6 +103,22 @@ def get_master_key(cls: Type["BaseKey"], passwd: bytes) -> bytes:
         """
         return get_master_key(cls.AUTH_ALG, passwd)
 
+    @classmethod
+    def get_localized_key(
+        cls: Type["BaseKey"], master_key: bytes, engine_id: bytes
+    ) -> bytes:
+        """
+        Convert master key to localized key.
+
+        Args:
+            master_key: Master key, must have size according to algorithm.
+            engine_id: SNMP engine id.
+
+        Returns:
+            Localized key. Resulting length same as master_key.
+        """
+        return get_localized_key(cls.AUTH_ALG, master_key, engine_id)
+
     def _pad(self: "BaseKey", key_len: int) -> None:
         """
         Pad key to given length.
diff --git a/src/lib.rs b/src/lib.rs
index 82c2b1a..d2aba4c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -31,5 +31,6 @@ fn gufo_ping(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<socket::GetNextIter>()?;
     m.add_class::<socket::GetBulkIter>()?;
     m.add_function(wrap_pyfunction!(util::get_master_key, m)?)?;
+    m.add_function(wrap_pyfunction!(util::get_localized_key, m)?)?;
     Ok(())
 }
diff --git a/src/util.rs b/src/util.rs
index 58b7c72..c02d500 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -6,6 +6,7 @@
 // ------------------------------------------------------------------------
 
 use crate::auth::{AuthKey, SnmpAuth};
+use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::types::PyBytes;
 
@@ -17,3 +18,21 @@ pub fn get_master_key(py: Python, alg: u8, passwd: &[u8]) -> PyResult<PyObject>
     auth.password_to_master(passwd, &mut out);
     Ok(PyBytes::new(py, &out).into())
 }
+
+// Convert master key to localized key
+#[pyfunction]
+pub fn get_localized_key(
+    py: Python,
+    alg: u8,
+    master_key: &[u8],
+    engine_id: &[u8],
+) -> PyResult<PyObject> {
+    let auth = AuthKey::new(alg)?;
+    let ks = auth.get_key_size();
+    if master_key.len() != ks {
+        return Err(PyValueError::new_err("invalid key size"));
+    }
+    let mut out = vec![0u8; auth.get_key_size()];
+    auth.localize(master_key, engine_id, &mut out);
+    Ok(PyBytes::new(py, &out).into())
+}
diff --git a/tests/test_user.py b/tests/test_user.py
index d959abf..9076c39 100644
--- a/tests/test_user.py
+++ b/tests/test_user.py
@@ -66,3 +66,27 @@ def test_get_master_key(
 ) -> None:
     r = auth.get_master_key(passwd)
     assert r == expected
+
+
+@pytest.mark.parametrize(
+    ("auth", "master_key", "engine_id", "expected"),
+    [
+        (
+            Md5Key,
+            b"\x9f\xaf\x32\x83\x88\x4e\x92\x83\x4e\xbc\x98\x47\xd8\xed\xd9\x63",
+            b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02",
+            b"\x52\x6f\x5e\xed\x9f\xcc\xe2\x6f\x89\x64\xc2\x93\x07\x87\xd8\x2b",
+        ),
+        (
+            Sha1Key,
+            b"\x9f\xb5\xcc\x03\x81\x49\x7b\x37\x93\x52\x89\x39\xff\x78\x8d\x5d\x79\x14\x52\x11",
+            b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02",
+            b"\x66\x95\xfe\xbc\x92\x88\xe3\x62\x82\x23\x5f\xc7\x15\x1f\x12\x84\x97\xb3\x8f\x3f",
+        ),
+    ],
+)
+def test_get_localized_key(
+    auth: BaseAuthKey, master_key: bytes, engine_id: bytes, expected: bytes
+) -> None:
+    r = auth.get_localized_key(master_key, engine_id)
+    assert r == expected