Plugin version comparison should work for all cases (#599)

Author: xen0l

aws_gate/decorators.py

@@ -2,6 +2,7 @@
 import os
 from subprocess import PIPE
 
+from packaging.version import parse as parse_version
 from wrapt import decorator
 
 from aws_gate.constants import PLUGIN_INSTALL_PATH, PLUGIN_NAME
@@ -43,9 +44,7 @@ def wrapper(
             required_version,
         )
 
-        if version and int(version.replace(".", "")) < int(
-            required_version.replace(".", "")
-        ):
+        if version and parse_version(version) < parse_version(required_version):
             raise ValueError("Invalid plugin version: {}".format(version))
 
         return wrapped_function(*args, **kwargs)

requirements/requirements.txt

@@ -1,7 +1,8 @@
 boto3~=1.15.15
+cryptography==3.1.1
 marshmallow==3.8.0
+packaging==20.4
 PyYAML>=5.1,<5.4
 requests==2.24.0
-cryptography==3.1.1
-wrapt==1.12.1
 unix-ar==0.2.1
+wrapt==1.12.1

tests/unit/test_decorators.py

@@ -41,8 +41,9 @@ def test_function():
         test_function()
 
 
-def test_plugin_version(mocker):
-    m = mocker.patch("aws_gate.decorators.execute_plugin", return_value="1.1.23.0")
+@pytest.mark.parametrize("version", ["1.1.23.0", "1.2.7.0"])
+def test_plugin_version(mocker, version):
+    m = mocker.patch("aws_gate.decorators.execute_plugin", return_value=version)
 
     @plugin_version("1.1.23.0")
     def test_function():
@@ -52,10 +53,11 @@ def test_function():
     assert m.call_args == mocker.call(["--version"], stdout=PIPE, stderr=PIPE)
 
 
-def test_plugin_version_bad_version(mocker):
+@pytest.mark.parametrize("version", ["1.1.25.0", "1.2.7.0"])
+def test_plugin_version_bad_version(mocker, version):
     mocker.patch("aws_gate.decorators.execute_plugin", return_value="1.1.23.0")
 
-    @plugin_version("1.1.25.0")
+    @plugin_version(version)
     def test_function():
         return "executed"