initial commit

Author: LiuG-lynx

README.md

@@ -1 +1,22 @@
-# postgresql-watcher
+# postgresql-watcher
+
+
+Casbin role watcher to be used for monitoring updates to casbin policies
+## Installation
+```bash
+pip install postgresql-watcher
+```
+
+## Basic Usage Example 
+### With Flask-authz
+```python
+from flask_authz import CasbinEnforcer
+from postgresql_watcher import PostgresqlWatcher
+from flask import Flask
+from casbin.persist.adapters import FileAdapter
+
+casbin_enforcer = CasbinEnforcer(app, adapter)
+watcher = PostgresqlWatcher(host=HOST, port=PORT, user=USER, password=PASSWORD)
+watcher.set_update_callback(casbin_enforcer.e.load_policy())
+casbin_enforcer.set_watcher(watcher)
+```

dev_requirements.txt

@@ -0,0 +1 @@
+black==20.8b1

examples/rbac_model.conf

@@ -0,0 +1,14 @@
+[request_definition]
+r = sub, obj, act
+
+[policy_definition]
+p = sub, obj, act
+
+[role_definition]
+g = _, _
+
+[policy_effect]
+e = some(where (p.eft == allow))
+
+[matchers]
+m = (p.sub == "*" || g(r.sub, p.sub)) && r.obj == p.obj && (p.act == "*" || r.act == p.act)

examples/rbac_policy.csv

@@ -0,0 +1,5 @@
+p,alice,data1,read
+p,bob,data2,write
+p,data2_admin,data2,read
+p,data2_admin,data2,write
+g,alice,data2_admin

postgresql_watcher/__init__.py

@@ -0,0 +1 @@
+from .watcher import PostgresqlWatcher

postgresql_watcher/watcher.py

@@ -0,0 +1,119 @@
+from typing import Optional, Callable, Any
+from psycopg2 import connect, extensions
+from multiprocessing import Process, Pipe, connection
+import time
+from select import select
+
+POSTGRESQL_CHANNEL_NAME = "casbin_role_watcher"
+
+
+def casbin_subscription(
+    process_conn: connection.PipeConnection,
+    host: str,
+    user: str,
+    password: str,
+    port: Optional[int] = 5432,
+    delay: Optional[int] = 2,
+    channel_name: Optional[str] = POSTGRESQL_CHANNEL_NAME,
+):
+    # delay connecting to postgresql (postgresql connection failure)
+    time.sleep(delay)
+    conn = connect(host=host, port=port, user=user, password=password)
+    # Can only receive notifications when not in transaction, set this for easier usage
+    conn.set_isolation_level(extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+    curs = conn.cursor()
+    curs.execute(f"LISTEN {channel_name};")
+    print("Waiting for casbin policy update")
+    while True and not curs.closed:
+        if not select([conn], [], [], 5) == ([], [], []):
+            print("Casbin policy update identified..")
+            conn.poll()
+            while conn.notifies:
+                notify = conn.notifies.pop(0)
+                print(f"Notify: {notify.payload}")
+                process_conn.put(notify.payload)
+
+
+
+class PostgresqlWatcher(object):
+    def __init__(
+        self,
+        host: str,
+        user: str,
+        password: str,
+        port: Optional[int] = 5432,
+        channel_name: Optional[str] = POSTGRESQL_CHANNEL_NAME,
+        start_process: Optional[bool] = True,
+    ):
+        self.host = host
+        self.port = port
+        self.user = user
+        self.password = password
+        self.channel_name = channel_name
+        self.subscribed_process, self.parent_conn = self.create_subscriber_process(
+            start_process
+        )
+
+    def create_subscriber_process(
+        self,
+        start_process: Optional[bool] = True,
+        delay: Optional[int] = 2,
+    ):
+        parent_conn, child_conn = Pipe()
+        p = Process(
+            target=casbin_subscription,
+            args=(
+                child_conn,
+                self.host,
+                self.user,
+                self.password,
+                self.port,
+                delay,
+                self.channel_name,
+            ),
+            daemon=True,
+        )
+        if start_process:
+            p.start()
+        self.should_reload()
+        return p, parent_conn
+
+    def update_callback(self):
+        print("callback called because casbin role updated")
+
+    def set_update_callback(self, fn_name: Any):
+        print("runtime is set update callback",fn_name)
+        self.update_callback = fn_name
+
+
+    def update(self):
+        conn = connect(
+            host=self.host,
+            port=self.port,
+            user=self.user,
+            password=self.password,
+        )
+        # Can only receive notifications when not in transaction, set this for easier usage
+        conn.set_isolation_level(extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+        curs = conn.cursor()
+        curs.execute(
+            f"NOTIFY {self.channel_name},'casbin policy update at {time.time()}'"
+        )
+        conn.close()
+        return True
+
+    def should_reload(self):
+        try:
+            if self.parent_conn.poll():
+                message = self.parent_conn.recv()
+                print(f"message:{message}")
+                return True
+        except EOFError:
+            print(
+                "Child casbin-watcher subscribe process has stopped, "
+                "attempting to recreate the process in 10 seconds..."
+            )
+            self.subscribed_process, self.parent_conn = self.create_subscriber_process(
+                delay=10
+            )
+            return False

requirements.txt

@@ -0,0 +1,2 @@
+casbin==0.8.4
+psycopg2-binary==2.8.6

tests/test_postgresql_watcher.py

@@ -0,0 +1,43 @@
+import unittest
+from postgresql_watcher import PostgresqlWatcher
+from multiprocessing import connection, context
+
+# Warning!!! , Please setup yourself config
+HOST = "127.0.0.1"
+PORT = 5432
+USER = "postgres"
+PASSWORD = "123456"
+
+
+def get_watcher():
+    pg_watcher = PostgresqlWatcher(host=HOST, port=PORT, user=USER, password=PASSWORD)
+    return pg_watcher
+
+
+pg_watcher = get_watcher()
+
+
+class TestConfig(unittest.TestCase):
+    def test_pg_watcher_init(self):
+        assert isinstance(pg_watcher.parent_conn, connection.PipeConnection)
+        assert isinstance(pg_watcher.subscribed_process, context.Process)
+
+    def test_update_pg_watcher(self):
+        assert pg_watcher.update() is True
+
+    def test_not_reload(self):
+        assert not pg_watcher.should_reload()
+
+    def test_default_update_callback(self):
+        assert pg_watcher.update_callback() is None
+
+    def test_add_update_callback(self):
+        def _test_callback():
+            pass
+
+        pg_watcher.set_update_callback(_test_callback)
+        assert pg_watcher.update_callback == _test_callback
+
+
+if __name__ == "__main__":
+    unittest.main()