Merge pull request #1 from grafolean/feature/history

Feature/history

Author: grafolean

.env.example

@@ -1,4 +1,5 @@
 BACKEND_URL=https://grafolean.com/api
 BOT_TOKEN=
+DB_DIR=
 JOBS_REFRESH_INTERVAL=60
 NETFLOW_PORT=2055

.gitignore

@@ -1,2 +1,4 @@
 .vscode
 .env
+.cache
+__pycache__

.gitmodules

@@ -1,3 +1,3 @@
 [submodule "pynetflow"]
 	path = pynetflow
-	url = https://github.com/bitkeks/python-netflow-v9-softflowd.git
+	url = https://github.com/grafolean/python-netflow-v9-softflowd.git

Dockerfile

@@ -9,8 +9,8 @@ FROM python:3.6-slim-stretch as build-backend
 COPY ./ /netflowbot/
 WORKDIR /netflowbot
 RUN \
+    rm -rf .git/ tests/ .vscode/ .pytest_cache/ __pycache__/ && \
     find ./ ! -name '*.py' -type f -exec rm '{}' ';' && \
-    rm -rf tests/ .vscode/ .pytest_cache/ __pycache__/ && \
     python3.6 -m compileall -b ./ && \
     find ./ -name '*.py' -exec rm '{}' ';'
 
@@ -38,8 +38,10 @@ WORKDIR /netflowbot
 HEALTHCHECK --interval=10s --retries=1 CMD /bin/bash -c "[ ! -f /tmp/fail_health_check ] || ( rm /tmp/fail_health_check && exit 1 )"
 
 # CAREFUL:
-# There are two entrypoints, both of which should be running:
-# - netflowcollector: gathering packets and writing statistics to Redis
+# There are three entrypoints, all of which should be running: (use docker-compose.yml to start 3 services)
+# - netflowcollector: gathering packets and writing them to named pipe
+# - netflowwriter: reading packets from named pipe and writing them to DB
 # - netflowbot: Grafolean bot for NetFlow - sending data to Grafolean according to configured sensors
-#CMD ["python", "-m", "netflowcollector"]
+# CMD ["python", "-m", "netflowcollector"]
+# CMD ["python", "-m", "netflowwriter"]
 CMD ["python", "-m", "netflowbot"]

Pipfile

@@ -11,7 +11,7 @@ requests = "*"
 python-dotenv = "*"
 ansicolors = "*"
 grafoleancollector = "*"
-redis = "*"
+psycopg2-binary = "*"
 
 [requires]
 python_version = "3.6"

Pipfile.lock

@@ -0,0 +1,157 @@
+from colors import color
+from contextlib import contextmanager
+import logging
+import os
+import sys
+import copy
+import json
+
+import psycopg2
+from psycopg2.pool import ThreadedConnectionPool
+from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT, register_adapter
+from psycopg2.extras import Json
+
+IS_DEBUG = os.environ.get('DEBUG', 'false') in ['true', 'yes', '1']
+logging.basicConfig(format='%(asctime)s.%(msecs)03d | %(levelname)s | %(message)s',
+                    datefmt='%Y-%m-%d %H:%M:%S', level=logging.DEBUG if IS_DEBUG else logging.INFO)
+logging.addLevelName(logging.DEBUG, color("DBG", 7))
+logging.addLevelName(logging.INFO, "INF")
+logging.addLevelName(logging.WARNING, color('WRN', fg='red'))
+logging.addLevelName(logging.ERROR, color('ERR', bg='red'))
+log = logging.getLogger("{}.{}".format(__name__, "dbutils"))
+
+
+db_pool = None
+DB_PREFIX = 'netflow_'
+register_adapter(dict, Json)
+
+
+# https://medium.com/@thegavrikstory/manage-raw-database-connection-pool-in-flask-b11e50cbad3
+@contextmanager
+def get_db_connection():
+    global db_pool
+    if db_pool is None:
+        db_connect()
+
+    try:
+        conn = db_pool.getconn()
+        conn.autocommit = True
+        conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
+        yield conn
+    finally:
+        db_pool.putconn(conn)
+
+
+@contextmanager
+def get_db_cursor(commit=False):
+    with get_db_connection() as connection:
+        cursor = connection.cursor()
+        try:
+            yield cursor
+            if commit:
+                connection.commit()
+        finally:
+            cursor.close()
+
+
+def db_connect():
+    global db_pool
+    host, dbname, user, password, connect_timeout = (
+        os.environ.get('DB_HOST', 'localhost'),
+        os.environ.get('DB_DATABASE', 'grafolean'),
+        os.environ.get('DB_USERNAME', 'admin'),
+        os.environ.get('DB_PASSWORD', 'admin'),
+        int(os.environ.get('DB_CONNECT_TIMEOUT', '10'))
+    )
+    try:
+        log.info("Connecting to database, host: [{}], db: [{}], user: [{}]".format(host, dbname, user))
+        db_pool = ThreadedConnectionPool(1, 20,
+                              database=dbname,
+                              user=user,
+                              password=password,
+                              host=host,
+                              port=5432,
+                              connect_timeout=connect_timeout)
+    except:
+        db_pool = None
+        log.error("DB connection failed")
+
+
+###########################
+#   DB schema migration   #
+###########################
+
+def get_existing_schema_version():
+    existing_schema_version = 0
+    with get_db_cursor() as c:
+        try:
+            c.execute(f'SELECT schema_version FROM {DB_PREFIX}runtime_data;')
+            res = c.fetchone()
+            existing_schema_version = res[0]
+        except psycopg2.ProgrammingError:
+            pass
+    return existing_schema_version
+
+
+def _get_migration_method(next_migration_version):
+    method_name = 'migration_step_{}'.format(next_migration_version)
+    return method_name if hasattr(sys.modules[__name__], method_name) else None
+
+
+def is_migration_needed():
+    existing_schema_version = get_existing_schema_version()
+    return _get_migration_method(existing_schema_version + 1) is not None
+
+
+def migrate_if_needed():
+    existing_schema_version = get_existing_schema_version()
+    try_migrating_to = existing_schema_version + 1
+    while True:
+        method_name = _get_migration_method(try_migrating_to)
+        if method_name is None:
+            break
+        log.info("Migrating DB schema from {} to {}".format(existing_schema_version, try_migrating_to))
+        method_to_call = getattr(sys.modules[__name__], method_name)
+        method_to_call()
+        # automatically upgrade schema version if there is no exception:
+        with get_db_cursor() as c:
+            c.execute(f'UPDATE {DB_PREFIX}runtime_data SET schema_version = %s;', (try_migrating_to,))
+        try_migrating_to += 1
+    if try_migrating_to == existing_schema_version + 1:
+        return False  # migration wasn't meeded
+    else:
+        return True
+
+
+def migration_step_1():
+    with get_db_cursor() as c:
+        c.execute(f'CREATE TABLE {DB_PREFIX}runtime_data (schema_version SMALLSERIAL NOT NULL);')
+        c.execute(f'INSERT INTO {DB_PREFIX}runtime_data (schema_version) VALUES (1);')
+
+def migration_step_2():
+    with get_db_cursor() as c:
+        # UNLOGGED: Disabling WAL avoids high I/O load. Since NetFlow data is of temporary nature, this still
+        # allows us to perform queries, but if the database crashes it is acceptable to lose all of the records.
+        c.execute(f'CREATE UNLOGGED TABLE {DB_PREFIX}records (seq BIGSERIAL NOT NULL PRIMARY KEY, ts NUMERIC(16,6) NOT NULL, client_ip TEXT);')
+
+        c.execute(f"""
+            CREATE UNLOGGED TABLE {DB_PREFIX}flows (
+                record INTEGER NOT NULL REFERENCES {DB_PREFIX}records(seq) ON DELETE CASCADE,
+                IN_BYTES INTEGER,
+                PROTOCOL SMALLINT,
+                DIRECTION SMALLINT,
+                L4_DST_PORT INTEGER,
+                L4_SRC_PORT INTEGER,
+                INPUT_SNMP SMALLINT,
+                OUTPUT_SNMP SMALLINT,
+                IPV4_DST_ADDR TEXT,
+                IPV4_SRC_ADDR TEXT
+            );
+        """)
+        c.execute(f'CREATE INDEX netflow_flows_record on netflow_flows (record);')
+
+        c.execute(f'CREATE TABLE {DB_PREFIX}bot_jobs (job_id TEXT NOT NULL PRIMARY KEY, last_used_seq BIGSERIAL);')
+
+        # slow db queries should be logged to help avoid performance problems:
+        database_name = os.environ.get('DB_DATABASE', 'grafolean')
+        c.execute(f'ALTER DATABASE %s SET log_min_duration_statement = 100;', (database_name,))