Use logger for all sorts of logging. Also prepare for writing some logs to a database for audit reasons.feature.scanners

5 files changed, 65 insertions, 78 deletions

diff --git a/src/db/dictionary.py b/src/db/dictionary.py
index f0f5fe9..bd468f1 100755
--- a/src/db/dictionary.py
+++ b/src/db/dictionary.py
@@ -11,9 +11,13 @@ import os
 import sys
 import time
 
+from log import get_logger
+
 from db import couch
 from db.schema import as_index_list, validate_collector_data
 
+logger = get_logger()
+
 
 class DictDB():
     def __init__(self):
@@ -42,9 +46,9 @@ class DictDB():
 
         try:
             self.couchdb = self.server.database(self.database)
-            print("Database already exists")
+            logger.debug("Database already exists")
         except couch.exceptions.NotFound:
-            print("Creating database and indexes.")
+            logger.debug("Creating database and indexes.")
             self.couchdb = self.server.create(self.database)
 
             for i in as_index_list():
diff --git a/src/db/index.py b/src/db/index.py
deleted file mode 100644
index 688ceeb..0000000
--- a/src/db/index.py
+++ /dev/null
@@ -1,61 +0,0 @@
-from pydantic import BaseSettings
-
-
-class CouchIindex(BaseSettings):
-    domain: dict = {
-        "index": {
-            "fields": [
-                "domain",
-            ]
-        },
-        "name": "domain-json-index",
-        "type": "json"
-    }
-    ip: dict = {
-        "index": {
-            "fields": [
-                "domain",
-                "ip"
-            ]
-        },
-        "name": "ip-json-index",
-        "type": "json"
-    }
-    port: dict = {
-        "index": {
-            "fields": [
-                "domain",
-                "port"
-            ]
-        },
-        "name": "port-json-index",
-        "type": "json"
-    }
-    asn: dict = {
-        "index": {
-            "fields": [
-                "domain",
-                "asn"
-            ]
-        },
-        "name": "asn-json-index",
-        "type": "json"
-    }
-    asn_country_code: dict = {
-        "index": {
-            "fields": [
-                "domain",
-                "asn_country_code"
-            ]
-        },
-        "name": "asn-country-code-json-index",
-        "type": "json"
-    }
-
-
-def as_list():
-    index_list = list()
-    for item in CouchIindex().dict():
-        index_list.append(CouchIindex().dict()[item])
-
-    return index_list
diff --git a/src/db/sql.py b/src/db/sql.py
index c47a69c..858ab6b 100644
--- a/src/db/sql.py
+++ b/src/db/sql.py
@@ -15,11 +15,12 @@ metadata = Base.metadata
 class Log(Base):
     __tablename__ = "log"
 
-    id = Column(Integer, primary_key=True)
-    timestamp = Column(Date, nullable=False,
-                       default=datetime.datetime.utcnow)
-    username = Column(Text, nullable=False)
-    logtext = Column(Text, nullable=False)
+    id = Column(Integer, primary_key=True, server_default=text(
+        "nextval('log_id_seq'::regclass)"))
+    name = Column(Text, nullable=False)
+    level = Column(Text, nullable=False)
+    msg = Column(Text, nullable=False)
+    func = Column(Text, nullable=True)
 
     def as_dict(self):
         """Return JSON serializable dict."""
@@ -34,12 +35,17 @@ class Log(Base):
         return d
 
     @classmethod
-    def add(cls, username, logtext):
-        with sqla_session() as session:
-            logentry = Log()
-            logentry.username = username
-            logentry.logtext = logtext
-            session.add(logentry)
+    def add(cls, name, level, msg, func=""):
+        try:
+            with sqla_session() as session:
+                logentry = Log()
+                logentry.name = name
+                logentry.level = level
+                logentry.msg = msg
+                logentry.func = func
+                session.add(logentry)
+        except Exception as e:
+            print(f"Failed to add logentry: {e}")
 
 
 class Scanner(Base):
diff --git a/src/log.py b/src/log.py
new file mode 100644
index 0000000..ec63106
--- /dev/null
+++ b/src/log.py
@@ -0,0 +1,26 @@
+import logging
+
+import uvicorn
+
+log_format = "%(levelprefix)s %(asctime)s | %(message)s"
+
+
+def get_handler():
+    formatter = uvicorn.logging.DefaultFormatter(log_format)
+    handler = logging.StreamHandler()
+    handler.setFormatter(formatter)
+
+    return handler
+
+
+def get_logger(name=None):
+    logger = logging.getLogger(name)
+    if logger.handlers:
+        return logger
+
+    logger = logging.getLogger(name)
+    handler = get_handler()
+    logger.addHandler(handler)
+    logger.setLevel(logging.DEBUG)
+
+    return logger
diff --git a/src/main.py b/src/main.py
index a62d77c..4a1741e 100755
--- a/src/main.py
+++ b/src/main.py
@@ -1,4 +1,5 @@
 import json
+import logging
 import os
 import sys
 import time
@@ -13,9 +14,12 @@ from pydantic import BaseModel
 
 from db.dictionary import DictDB
 from db.schema import get_index_keys
+from db.sql import Scanner
+from log import get_logger, log_format
 
-app = FastAPI()
+logger = get_logger()
 
+app = FastAPI()
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["http://localhost:8000"],
@@ -38,12 +42,12 @@ for i in range(10):
     try:
         db = DictDB()
     except Exception as e:
-        print(f"Database not responding, will try again soon: {e}")
+        logger.info(f"Database not responding, will try again soon: {e}")
     else:
         break
     time.sleep(1)
 else:
-    print('Database did not respond after 10 attempts, quitting.')
+    logger.info('Database did not respond after 10 attempts, quitting.')
     sys.exit(-1)
 
 
@@ -57,7 +61,7 @@ def get_pubkey():
         with open(keypath, "r") as fd:
             pubkey = fd.read()
     except FileNotFoundError:
-        print(f"Could not find JWT certificate in {keypath}")
+        logger.info(f"Could not find JWT certificate in {keypath}")
         sys.exit(-1)
 
     return pubkey
@@ -191,6 +195,9 @@ async def add(data: Request, Authorize: AuthJWT = Depends()):
             status_code=400,
         )
 
+    record_id = key["_id"]
+    logger.info(f"Record {record_id} added")
+
     return JSONResponse(content={"status": "success", "docs": key})
 
 
@@ -257,6 +264,11 @@ async def scanner_put(name, data: Request, Authorize: AuthJWT = Depends()):
 
 
 def main(standalone=False):
+    log_config = uvicorn.config.LOGGING_CONFIG
+    log_config["formatters"]["access"]["fmt"] = log_format
+    log_config["formatters"]["default"]["fmt"] = log_format
+    log_config["loggers"]["uvicorn"]["propagate"] = False
+
     if not standalone:
         return app