Bug fixes and added tests

Author: kartikbehl99

app.py

@@ -13,3 +13,7 @@ def shutdown_session() -> None:
     """
 
     db_session.remove()
+
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=5000, debug=True, threaded=True)

ssh_manager_backend/app/models/access_control.py

@@ -77,12 +77,15 @@ def grant_access(self, username: str, ip_addresses: List[str]) -> bool:
 
         return True
 
-    def revoke_access(self, username: str, ip_addresses: List[str]) -> bool:
+    def revoke_access(
+        self, username: str, ip_addresses: List[str], revoke_all: bool = False
+    ) -> bool:
         """
         Updates user access.
 
         :param username:
         :param ip_addresses:
+        :param revoke_all:.
         :return: booleans value for success/failure.
         """
 
@@ -91,15 +94,21 @@ def revoke_access(self, username: str, ip_addresses: List[str]) -> bool:
                 AccessControl.username == username
             ).first()
 
-            for ip in ip_addresses:
-                try:
-                    acl_details.ip_addresses.remove(ip)
-                except ValueError:
-                    continue
+            if not revoke_all:
+                for ip in ip_addresses:
+                    try:
+                        acl_details.ip_addresses.remove(ip)
+                    except ValueError:
+                        continue
+
+                self.session.query(AccessControl).filter(
+                    AccessControl.username == username
+                ).update({"ip_addresses": acl_details.ip_addresses})
+            else:
+                self.session.query(AccessControl).filter(
+                    AccessControl.username == username
+                ).update({"ip_addresses": []})
 
-            self.session.query(AccessControl).filter(
-                AccessControl.username == username
-            ).update({"ip_addresses": acl_details.ip_addresses})
             self.session.commit()
         except AttributeError:
             return False

ssh_manager_backend/app/models/keys.py

@@ -21,26 +21,30 @@ def exists(self, key_name: str) -> bool:
         return self.session.query(Key).filter(Key.name == key_name).first() is not None
 
     def create(
-        self, name: str, encrypted_key: bytes, key_hash: str, user: User
+        self, name: str, encrypted_key: bytes, key_hash: str, user_id: int
     ) -> bool:
         """
         Creates a key in database.
 
         :param name:
         :param encrypted_key:
         :param key_hash:
-        :param user:
+        :param user_id:
         :return: Boolean value indicating success/failure.
         """
 
         try:
             key: Key = Key(
-                name=name, encrypted_key=encrypted_key, key_hash=key_hash, user=user
+                name=name,
+                encrypted_key=encrypted_key,
+                key_hash=key_hash,
+                user_id=user_id,
             )
 
             self.session.add(key)
             self.session.commit()
         except SQLAlchemyError:
+            self.session.rollback()
             return False
 
         return True

ssh_manager_backend/app/models/keys_mapping.py

@@ -25,21 +25,21 @@ def exists(self, ip_address: str) -> bool:
             is not None
         )
 
-    def create(self, ip_address: str, key_name: str, key: Key) -> bool:
+    def create(self, ip_address: str, key_name: str) -> bool:
         """
         Creates a key mapping in db.
 
         :param ip_address:
         :param key_name:
-        :param key: Key object
         :return: Boolean value indicating success/failure.
         """
 
         try:
-            key_mapping = KeyMapping(key_name=key_name, ip_address=ip_address, key=key)
+            key_mapping = KeyMapping(key_name=key_name, ip_address=ip_address)
             self.session.add(key_mapping)
             self.session.commit()
         except SQLAlchemyError:
+            self.session.rollback()
             return False
 
         return True

ssh_manager_backend/db/database.py

@@ -3,7 +3,8 @@
 from sqlalchemy.orm import scoped_session, sessionmaker
 
 DB_URI = (
-    "postgresql+psycopg2://postgres:vP28ObNJLhb5qFDe@35.222.241.198/ssh_manager_test"
+    # "postgresql+psycopg2://postgres:vP28ObNJLhb5qFDe@35.222.241.198/ssh_manager_test"
+    "postgresql+psycopg2://ssh_manager:pass@localhost/ssh_manager_dev"
 )
 engine = create_engine(DB_URI, echo=True)
 db_session = scoped_session(

tests/models/key_mapping_model_test.py

@@ -0,0 +1,96 @@
+import pytest
+
+from ssh_manager_backend.app.models.keys import KeyModel
+from ssh_manager_backend.app.models.keys_mapping import KeyMappingModel
+from ssh_manager_backend.app.models.user import UserModel
+from tests.test_ssh_manager_backend import db_cleanup
+
+
+class TestKeyMappingModel:
+    @pytest.fixture
+    def cleanup(self):
+        yield
+        db_cleanup()
+
+    def test_create(self):
+        user = UserModel()
+        name: str = "test_user"
+        username: str = "test_username"
+        password: str = "test_password"
+        admin: bool = False
+        encrypted_dek: bytes = b"test_encrypted_dek"
+        iv_for_dek: bytes = b"test_iv_for_dek"
+        salt_for_dek: bytes = b"test_salt_for_dek"
+        iv_for_kek: bytes = b"test_iv_for_kek"
+        salt_for_kek: bytes = b"test_salt_for_kek"
+        salt_for_password: bytes = b"test_salt_for_password"
+
+        key = KeyModel()
+        key_name: str = "test_key"
+        encrypted_key: bytes = b"encrypted_test_key"
+        key_hash: str = "test_key_hash"
+
+        assert (
+            key.create(
+                name=key_name,
+                encrypted_key=encrypted_key,
+                key_hash=key_hash,
+                user_id=12345,
+            )
+            is False
+        )
+
+        assert (
+            user.create(
+                name=name,
+                username=username,
+                password=password,
+                admin=admin,
+                encrypted_dek=encrypted_dek,
+                iv_for_dek=iv_for_dek,
+                salt_for_dek=salt_for_dek,
+                iv_for_kek=iv_for_kek,
+                salt_for_kek=salt_for_kek,
+                salt_for_password=salt_for_password,
+            )
+            is True
+        )
+
+        user_id: int = user.get_user(username=username).id
+
+        assert (
+            key.create(
+                name=key_name,
+                encrypted_key=encrypted_key,
+                key_hash=key_hash,
+                user_id=user_id,
+            )
+            is True
+        )
+
+        key_mapping: KeyMappingModel = KeyMappingModel()
+        ip_address: str = "1.1.1.1"
+
+        assert key_mapping.create(ip_address=ip_address, key_name=key_name) is True
+
+        assert (
+            key_mapping.create(ip_address=ip_address, key_name="non_existent_key")
+            is False
+        )
+
+    def test_exists(self):
+        key_mapping: KeyMappingModel = KeyMappingModel()
+        ip_address: str = "1.1.1.1"
+
+        assert key_mapping.exists(ip_address=ip_address) is True
+
+        assert key_mapping.exists(ip_address="2.1.1.1") is False
+
+    def test_get_mapping(self, cleanup):
+        key_mapping: KeyMappingModel = KeyMappingModel()
+        key_name: str = "test_key"
+        ip_address: str = "1.1.1.1"
+
+        assert key_mapping.get_mapping(ip_address=ip_address).key_name == key_name
+
+        assert key_mapping.get_mapping(ip_address="2.2.2.2") is None

tests/models/key_model_test.py

@@ -0,0 +1,84 @@
+import pytest
+
+from ssh_manager_backend.app.models.keys import KeyModel
+from ssh_manager_backend.app.models.user import UserModel
+from tests.test_ssh_manager_backend import db_cleanup
+
+
+class TestKeyModel:
+    @pytest.fixture
+    def cleanup(self):
+        yield
+        db_cleanup()
+
+    def test_create(self):
+        user = UserModel()
+        name: str = "test_user"
+        username: str = "test_username"
+        password: str = "test_password"
+        admin: bool = False
+        encrypted_dek: bytes = b"test_encrypted_dek"
+        iv_for_dek: bytes = b"test_iv_for_dek"
+        salt_for_dek: bytes = b"test_salt_for_dek"
+        iv_for_kek: bytes = b"test_iv_for_kek"
+        salt_for_kek: bytes = b"test_salt_for_kek"
+        salt_for_password: bytes = b"test_salt_for_password"
+
+        key = KeyModel()
+        key_name: str = "test_key"
+        encrypted_key: bytes = b"encrypted_test_key"
+        key_hash: str = "test_key_hash"
+
+        assert (
+            key.create(
+                name=key_name,
+                encrypted_key=encrypted_key,
+                key_hash=key_hash,
+                user_id=12345,
+            )
+            is False
+        )
+
+        assert (
+            user.create(
+                name=name,
+                username=username,
+                password=password,
+                admin=admin,
+                encrypted_dek=encrypted_dek,
+                iv_for_dek=iv_for_dek,
+                salt_for_dek=salt_for_dek,
+                iv_for_kek=iv_for_kek,
+                salt_for_kek=salt_for_kek,
+                salt_for_password=salt_for_password,
+            )
+            is True
+        )
+
+        user_id = user.get_user(username=username).id
+
+        assert (
+            key.create(
+                name=key_name,
+                encrypted_key=encrypted_key,
+                key_hash=key_hash,
+                user_id=user_id,
+            )
+            is True
+        )
+
+    def test_exists(self):
+        key = KeyModel()
+        key_name: str = "test_key"
+
+        assert key.exists(key_name=key_name) is True
+
+        assert key.exists(key_name="non_existent_key") is False
+
+    def test_get_key(self, cleanup):
+        key = KeyModel()
+        key_name: str = "test_key"
+
+        assert key.get_key(key_name=key_name).name == key_name
+
+        assert key.get_key(key_name="non_existent_keyname") is None