Merge pull request #4 from thephilomaths/add-tests

Bug fix and added user model test

Author: kartikbehl99

.isort.cfg

@@ -2,4 +2,4 @@
 line_length = 88
 multi_line_output = 3
 include_trailing_comma = True
-known_third_party = cryptography,flask,sqlalchemy
+known_third_party = cryptography,flask,pytest,sqlalchemy

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)

requirements.txt

@@ -16,3 +16,30 @@ SQLAlchemy==1.3.19
 toml==0.10.1
 virtualenv==20.0.31
 Werkzeug==1.0.1
+appdirs==1.4.4
+attrs==20.1.0
+cfgv==3.2.0
+click==7.1.2
+distlib==0.3.1
+filelock==3.0.12
+Flask==1.1.2
+identify==1.4.29
+iniconfig==1.0.1
+itsdangerous==1.1.0
+Jinja2==2.11.2
+MarkupSafe==1.1.1
+more-itertools==8.5.0
+nodeenv==1.5.0
+packaging==20.4
+pluggy==0.13.1
+pre-commit==2.7.1
+psycopg2==2.8.5
+py==1.9.0
+pyparsing==2.4.7
+pytest==6.0.1
+PyYAML==5.3.1
+six==1.15.0
+SQLAlchemy==1.3.19
+toml==0.10.1
+virtualenv==20.0.31
+Werkzeug==1.0.1

ssh_manager_backend/app/models/access_control.py

@@ -10,6 +10,26 @@ class AccessControlModel:
     def __init__(self):
         self.session = db_session()
 
+    def create(self, username: str):
+        """
+        Creates an entry in the access_control table for the given user.
+
+        :param username
+        :return: boolean value whether the entry is created or not.
+        """
+
+        try:
+            acl_details: AccessControl = AccessControl(
+                username=username, ip_addresses=[]
+            )
+            self.session.add(acl_details)
+            self.session.commit()
+        except SQLAlchemyError:
+            self.session.rollback()
+            return False
+
+        return True
+
     def has_access(self, username: str, ip_address: str) -> bool:
         """
         Checks whether a user has access to the provided the list of ip addresses.
@@ -43,18 +63,29 @@ def grant_access(self, username: str, ip_addresses: List[str]) -> bool:
 
             acl_details.ip_addresses += ip_addresses
             acl_details.ip_addresses = list(set(acl_details.ip_addresses))
+
+            self.session.query(AccessControl).filter(
+                AccessControl.username == username
+            ).update({"ip_addresses": acl_details.ip_addresses})
+
             self.session.commit()
-        except [AttributeError, SQLAlchemyError]:
+        except AttributeError:
+            return False
+        except SQLAlchemyError:
+            self.session.rollback()
             return False
 
         return True
 
-    def remove_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.
         """
 
@@ -63,11 +94,26 @@ def remove_access(self, username: str, ip_addresses: List[str]) -> bool:
                 AccessControl.username == username
             ).first()
 
-            for ip in ip_addresses:
-                acl_details.ip_addresses.remove(ip)
+            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.commit()
-        except [AttributeError, SQLAlchemyError]:
+        except AttributeError:
+            return False
+        except SQLAlchemyError:
+            self.session.rollback()
             return False
 
         return True
@@ -85,5 +131,5 @@ def get_all_ips(self, username: str) -> List[str]:
                 AccessControl.username == username
             ).first()
             return acl_details.ip_addresses
-        except [AttributeError, SQLAlchemyError]:
+        except (AttributeError, SQLAlchemyError):
             return []

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/app/models/user.py

@@ -94,6 +94,7 @@ def create(
             self.session.add(user)
             self.session.commit()
         except SQLAlchemyError:
+            self.session.rollback()
             return False
 
         return True
@@ -124,3 +125,23 @@ def get_user(self, username: str) -> Union[None, User]:
         """
 
         return self.session.query(User).filter(User.username == username).first()
+
+    def destroy_user(self, username: str) -> bool:
+        """
+        Deletes a user.
+
+        :param username:
+        :return: Boolean value indicating success/failure.
+        """
+
+        try:
+            user: User = self.session.query(User).filter(
+                User.username == username
+            ).first()
+            self.session.delete(user)
+            self.session.commit()
+        except SQLAlchemyError:
+            self.session.rollback()
+            return False
+
+        return True

ssh_manager_backend/db/database.py

@@ -2,8 +2,11 @@
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import scoped_session, sessionmaker
 
-DB_URI = "postgresql+psycopg2://postgres:ssh_manager@localhost/ssh_manager_dev"
-engine = create_engine(DB_URI, convert_unicode=True, echo=True)
+DB_URI = (
+    "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(
     sessionmaker(autocommit=False, autoflush=False, bind=engine)
 )

ssh_manager_backend/db/schema.py

@@ -1,4 +1,4 @@
-from sqlalchemy import ARRAY, Binary, Boolean, Column, ForeignKey, Integer, String
+from sqlalchemy import ARRAY, Boolean, Column, ForeignKey, Integer, LargeBinary, String
 from sqlalchemy.orm import relationship
 
 from ssh_manager_backend.db.database import Base
@@ -12,12 +12,12 @@ class User(Base):
     username = Column(String, unique=True)
     password = Column(String)
     admin = Column(Boolean)
-    encrypted_dek = Column(Binary, unique=True)
-    iv_for_dek = Column(Binary, unique=True)
-    salt_for_dek = Column(Binary, unique=True)
-    iv_for_kek = Column(Binary, unique=True)
-    salt_for_kek = Column(Binary, unique=True)
-    salt_for_password = Column(Binary, unique=True)
+    encrypted_dek = Column(LargeBinary, unique=True)
+    iv_for_dek = Column(LargeBinary, unique=True)
+    salt_for_dek = Column(LargeBinary, unique=True)
+    iv_for_kek = Column(LargeBinary, unique=True)
+    salt_for_kek = Column(LargeBinary, unique=True)
+    salt_for_password = Column(LargeBinary, unique=True)
     keys = relationship("Key", backref="users")
     access_control = relationship("AccessControl", backref="users")
 
@@ -34,7 +34,7 @@ class Key(Base):
 
     id = Column(Integer, primary_key=True)
     name = Column(String, unique=True)
-    encrypted_key = Column(Binary, unique=True)
+    encrypted_key = Column(LargeBinary, unique=True)
     key_hash = Column(String, unique=True)
     user_id = Column(Integer, ForeignKey("users.id"))
     user = relationship("User")

tests/models/acl_model_test.py

@@ -0,0 +1,88 @@
+from typing import List
+
+import pytest
+
+from ssh_manager_backend.app.models.access_control import AccessControlModel
+from ssh_manager_backend.app.models.user import UserModel
+from tests.test_ssh_manager_backend import db_cleanup
+
+
+class TestAccessControlModel:
+    @pytest.fixture
+    def cleanup(self):
+        yield
+        db_cleanup()
+
+    def test_create(self):
+        acl: AccessControlModel = AccessControlModel()
+        user: UserModel = UserModel()
+
+        name: str = "test_user"
+        username = "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"
+
+        assert acl.create(username=username) 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
+        )
+
+        assert acl.create(username=username) is True
+
+    def test_grant_access(self):
+        acl: AccessControlModel = AccessControlModel()
+        username: str = "test_username"
+        ip_addresses: List[str] = ["1.1.1.1", "1.0.0.1"]
+
+        assert acl.grant_access(username=username, ip_addresses=ip_addresses) is True
+
+        assert (
+            acl.grant_access(
+                username="non_existent_username", ip_addresses=ip_addresses
+            )
+            is False
+        )
+
+        assert sorted(acl.get_all_ips(username=username)) == sorted(ip_addresses)
+
+    def test_revoke_access(self, cleanup):
+        acl: AccessControlModel = AccessControlModel()
+        username: str = "test_username"
+        ip_addresses: List[str] = ["1.1.1.1", "1.0.0.1"]
+
+        assert (
+            acl.revoke_access(username=username, ip_addresses=[ip_addresses[0]]) is True
+        )
+
+        assert (
+            acl.revoke_access(username=username, ip_addresses=["non_existent_ip"])
+            is True
+        )
+
+        assert (
+            acl.revoke_access(
+                username="non_existent_username", ip_addresses=ip_addresses
+            )
+            is False
+        )
+
+        assert acl.get_all_ips(username=username) == [ip_addresses[1]]