From 08b174926c7bb2540d79f53e9c21c1582d53d585 Mon Sep 17 00:00:00 2001
From: 13ph03nix <17541483+13ph03nix@users.noreply.github.com>
Date: Wed, 26 Jul 2023 12:59:43 -0700
Subject: [PATCH] refactor: making mmh3 an optional dependency

---
 .../protocols/common/expressions/__init__.py  | 22 ++++++++++++++-----
 requirements.txt                              |  1 -
 setup.py                                      |  4 ++--
 tests/test_nuclei_helper_functions.py         |  4 ++--
 4 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/pocsuite3/lib/yaml/nuclei/protocols/common/expressions/__init__.py b/pocsuite3/lib/yaml/nuclei/protocols/common/expressions/__init__.py
index 2c535172..c9afd97f 100644
--- a/pocsuite3/lib/yaml/nuclei/protocols/common/expressions/__init__.py
+++ b/pocsuite3/lib/yaml/nuclei/protocols/common/expressions/__init__.py
@@ -16,7 +16,6 @@
 from typing import get_type_hints, Union
 
 import chardet
-import mmh3 as py_mmh3
 from pkg_resources import parse_version
 from pocsuite3.lib.core.log import LOGGER as logger
 from pocsuite3.lib.yaml.nuclei.protocols.common.expressions.safe_eval import safe_eval
@@ -120,9 +119,11 @@ def base64_py(src: Union[bytes, str]) -> str:
 
     Example:
         Input: base64_py("Hello")
-        Output: SGVsbG8=
+        Output: "SGVsbG8=\n"
     """
-    return base64(src)
+    if not isinstance(src, bytes):
+        src = src.encode('utf-8')
+    return py_built_in_base64.encodebytes(src).decode('utf-8')
 
 
 @auto_convert_types
@@ -373,15 +374,24 @@ def md5(inp: Union[str, bytes]) -> str:
 
 
 @auto_convert_types
-def mmh3(inp: str) -> int:
+def mmh3(inp: str) -> str:
     """
     Calculates the MMH3 (MurmurHash3) hash of an input
 
     Example:
         Input: mmh3("Hello")
-        Output: 316307400
+        Output: "316307400"
     """
-    return py_mmh3.hash(inp)
+
+    try:
+        import mmh3 as py_mmh3
+    except ImportError:
+        logger.error('Python extension for MurmurHash (MurmurHash3) is not installed. '
+                     'Reason: https://github.com/knownsec/pocsuite3/issues/359, '
+                     'You can locate the packages here: https://pypi.org/project/mmh3/')
+        return "0"
+
+    return str(py_mmh3.hash(inp))
 
 
 def print_debug(*args) -> None:
diff --git a/requirements.txt b/requirements.txt
index 679d11c0..d8562b3d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -14,4 +14,3 @@ pycryptodomex >= 3.9.0
 dacite >= 1.6.0
 PyYAML >= 6.0
 lxml >= 4.6.0
-mmh3 >= 3.0.0
diff --git a/setup.py b/setup.py
index d8afe8a4..7764fafa 100644
--- a/setup.py
+++ b/setup.py
@@ -57,12 +57,12 @@ def find_packages(where='.'):
         "dacite",
         "PyYAML",
         "lxml",
-        "mmh3"
     ],
     extras_require={
         'complete': [
             'pyOpenSSL',
-            'jq'
+            'jq',
+            'mmh3'
         ],
     }
 )
diff --git a/tests/test_nuclei_helper_functions.py b/tests/test_nuclei_helper_functions.py
index 00cd54ae..9bfea58a 100644
--- a/tests/test_nuclei_helper_functions.py
+++ b/tests/test_nuclei_helper_functions.py
@@ -10,7 +10,7 @@ def test_base64_decode(self):
         self.assertEqual(base64_decode("SGVsbG8="), b"Hello")
 
     def test_base64_py(self):
-        self.assertEqual(base64_py("Hello"), "SGVsbG8=")
+        self.assertEqual(base64_py("Hello"), "SGVsbG8=\n")
 
     def test_concat(self):
         self.assertEqual(concat("Hello", 123, "world"), "Hello123world")
@@ -74,7 +74,7 @@ def test_md5(self):
         self.assertEqual(md5("Hello"), "8b1a9953c4611296a827abf8c47804d7")
 
     def test_mmh3(self):
-        self.assertEqual(mmh3("Hello"), 316307400)
+        self.assertEqual(mmh3("Hello"), "316307400")
 
     def test_rand_base(self):
         self.assertRegex(rand_base(5, "abc"), r"[abc]{5}")