Update config loading/validation

7 files changed, 106 insertions, 203 deletions

diff --git a/.gitignore b/.gitignore
index 640f67f..e9ec5c8 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,6 +3,6 @@ data.db
 __pycache__
 .DS_Store
 internal_notes.txt
-config.ini
+config.yaml
 data
 docker-volume
\ No newline at end of file
diff --git a/app/main.py b/app/main.py
index c58a45f..266c589 100644
--- a/app/main.py
+++ b/app/main.py
@@ -15,7 +15,7 @@ import random
 from models import User, Link
 from database import *
 from app.util.log import log
-from var import BASE_URL
+from config import BASE_URL
 
 
 class FlaskUser(UserMixin):
diff --git a/app/util/log.py b/app/util/log.py
index 2b3542b..943863d 100644
--- a/app/util/log.py
+++ b/app/util/log.py
@@ -4,7 +4,7 @@ from ua_parser import user_agent_parser
 from ip2locationio.ipgeolocation import IP2LocationIOAPIError
 
 from database import SessionLocal
-from var import LOG, API_KEY, IP_TO_LOCATION
+from config import LOG, API_KEY, IP_TO_LOCATION
 from models import Link, Record
 
 configuration = ip2locationio.Configuration(API_KEY)
diff --git a/config.py b/config.py
new file mode 100644
index 0000000..9c56add
--- /dev/null
+++ b/config.py
@@ -0,0 +1,101 @@
+import jsonschema
+import os
+import yaml
+import validators
+import sys
+import logging
+from colorlog import ColoredFormatter
+
+log_level = logging.DEBUG
+log_format = (
+    "  %(log_color)s%(levelname)-8s%(reset)s |"
+    " %(log_color)s%(message)s%(reset)s"
+)
+
+logging.root.setLevel(log_level)
+formatter = ColoredFormatter(log_format)
+
+stream = logging.StreamHandler()
+stream.setLevel(log_level)
+stream.setFormatter(formatter)
+
+LOG = logging.getLogger("pythonConfig")
+LOG.setLevel(log_level)
+LOG.addHandler(stream)
+
+BASE_URL = None
+IP_TO_LOCATION = None
+API_KEY = None
+
+schema = {
+    "type": "object",
+    "properties": {
+        "config": {
+            "type": "object",
+            "properties": {
+                "base_url": {"type": "string"},
+                "ip_to_location": {"type": "boolean"},
+                "api_key": {"type": "string"},
+            },
+            "required": ["base_url", "ip_to_location"],
+        }
+    },
+    "required": ["config"],
+}
+
+
+# Load config file or create new template
+def load_config():
+    if os.path.exists("/.dockerenv"):
+        file_path = "/data/config.yaml"
+    else:
+        file_path = "config.yaml"
+
+    try:
+        with open(file_path, "r") as f:
+            file_contents = f.read()
+            validate_config(file_contents)
+
+    except FileNotFoundError:
+        # Create new config.yaml w/ template
+        with open(file_path, "w") as f:
+            f.write(
+                """
+            base_url: ""
+            ip_to_location: ""
+            api_key: ""
+            """
+            )
+        LOG.critical(
+            "`config.yaml` was not found, a template has been created."
+            " Please fill out the necessary information and restart."
+        )
+        sys.exit()
+
+
+# Validate the options within config.yaml
+def validate_config(file_contents):
+    global BASE_URL, IP_TO_LOCATION, API_KEY
+    config = yaml.safe_load(file_contents)
+
+    try:
+        jsonschema.validate(config, schema)
+    except jsonschema.ValidationError as e:
+        LOG.error(e.message)
+        sys.exit()
+
+    # Validate BASE_URL
+    if not validators.url(config["config"]["base_url"]):
+        LOG.error("BASE_URL is not a valid URL")
+    else:
+        BASE_URL = config["config"]["base_url"]
+
+    # Make IP_TO_LOCATION a boolean
+    IP_TO_LOCATION = bool(config["config"]["ip_to_location"])
+
+    # Validate API_KEY if IP_TO_LOCATION is set to TRUE
+    if IP_TO_LOCATION:
+        if not config["config"]["api_key"]:
+            LOG.error("API_KEY is not set")
+        else:
+            API_KEY = config["config"]["api_key"]
diff --git a/linklogger.py b/linklogger.py
index 05f6c23..bc735a1 100644
--- a/linklogger.py
+++ b/linklogger.py
@@ -1,7 +1,7 @@
 from werkzeug.middleware.dispatcher import DispatcherMiddleware
 from a2wsgi import ASGIMiddleware
 
-from validate_config import validate_config
+from config import load_config
 from app.main import app as flask_app
 from api.main import app as fastapi_app
 from database import Base, engine
@@ -17,5 +17,5 @@ flask_app.wsgi_app = DispatcherMiddleware(
 )
 
 if __name__ == "__main__":
-    validate_config()
+    load_config()
     flask_app.run(port=5252)
diff --git a/validate_config.py b/validate_config.py
deleted file mode 100644
index f7e6589..0000000
--- a/validate_config.py
+++ /dev/null
@@ -1,139 +0,0 @@
-import configparser
-import validators
-import os
-import sys
-
-from var import LOG
-
-"""
-Validate the config of a Docker run (environment variables)
-"""
-
-
-def validate_docker_config():
-    errors = 0
-
-    # Validate BASE_URL
-    try:
-        if not os.environ["BASE_URL"]:
-            LOG.error("BASE_URL is not set")
-            errors += 1
-        elif not validators.url(os.environ["BASE_URL"]):
-            LOG.error("BASE_URL is not a valid URL")
-            errors += 1
-    except KeyError:
-        LOG.critical("BASE_URL does not exist!")
-        errors += 1
-
-    # Validate IP_TO_LOCATION
-    try:
-        if not os.environ["IP_TO_LOCATION"]:
-            LOG.error("IP_TO_LOCATION is not set")
-            errors += 1
-        elif os.environ["IP_TO_LOCATION"].upper() not in ["TRUE", "FALSE", "T", "F"]:
-            LOG.error("IP_TO_LOCATION is not set to TRUE or FALSE")
-            errors += 1
-        else:
-            iptolocation = (
-                True if os.environ["IP_TO_LOCATION"].upper() in ["TRUE", "T"] else False
-            )
-            # Validate API_KEY if IP_TO_LOCATION is set to TRUE
-            if iptolocation:
-                try:
-                    if not os.environ["API_KEY"]:
-                        LOG.error("API_KEY is not set")
-                        errors += 1
-                except KeyError:
-                    LOG.critical("API_KEY does not exist!")
-                    errors += 1
-    except KeyError:
-        LOG.critical("IP_TO_LOCATION does not exist!")
-        errors += 1
-
-    if errors > 0:
-        LOG.critical(f"{errors} error(s) found in environment variables")
-        sys.exit()
-
-
-"""
-Validate the config of a bare metal run (config.ini file)
-"""
-
-
-def validate_bare_metal_config(file_contents):
-
-    config = configparser.ConfigParser()
-    config.read_string(file_contents)
-
-    errors = 0
-
-    # Validate BASE_URL
-    try:
-        if not config["CONFIG"]["BASE_URL"]:
-            LOG.error("BASE_URL is not set")
-            errors += 1
-        elif not validators.url(config["CONFIG"]["BASE_URL"]):
-            LOG.error("BASE_URL is not a valid URL")
-            errors += 1
-    except ValueError:
-        LOG.critical("BASE_URL does not exist!")
-        errors += 1
-
-    # Validate IP_TO_LOCATION
-    try:
-        if not config["CONFIG"]["IP_TO_LOCATION"]:
-            LOG.error("IP_TO_LOCATION is not set")
-            errors += 1
-        elif config["CONFIG"]["IP_TO_LOCATION"].upper() not in [
-            "TRUE",
-            "FALSE",
-            "T",
-            "F",
-        ]:
-            LOG.error("IP_TO_LOCATION is not set to TRUE or FALSE")
-            errors += 1
-        else:
-            iptolocation = (
-                True
-                if config["CONFIG"]["IP_TO_LOCATION"].upper() in ["TRUE", "T"]
-                else False
-            )
-        # Validate API_KEY if IP_TO_LOCATION is set to TRUE
-        if iptolocation:
-            try:
-                if not config["CONFIG"]["API_KEY"]:
-                    LOG.error("API_KEY is not set")
-                    errors += 1
-            except ValueError:
-                LOG.critical("API_KEY does not exist!")
-                errors += 1
-    except ValueError:
-        LOG.critical("IP_TO_LOCATION does not exist!")
-        errors += 1
-
-    if errors > 0:
-        LOG.critical(f"{errors} error(s) found in `config.ini`")
-        sys.exit()
-
-
-def validate_config():
-    # If the app is running in Docker
-    if "BASE_URL" in os.environ or "IP_TO_LOCATION" in os.environ:
-        return validate_docker_config()
-
-    # Otherwise, the app is running on bare metal
-    try:
-        with open("config.ini", "r") as f:
-            file_contents = f.read()
-            return validate_bare_metal_config(file_contents)
-    except FileNotFoundError:
-        config = configparser.ConfigParser()
-        config["CONFIG"] = {"BASE_URL": "", "IP_TO_LOCATION": "", "API_KEY": ""}
-
-        with open("config.ini", "w") as configfile:
-            config.write(configfile)
-
-        LOG.error(
-            "`config.ini` has been created. Fill out the necessary information then re-run."
-        )
-        sys.exit()
diff --git a/var.py b/var.py
deleted file mode 100644
index e444629..0000000
--- a/var.py
+++ /dev/null
@@ -1,59 +0,0 @@
-import configparser
-import logging
-import os
-from colorlog import ColoredFormatter
-
-
-log_level = logging.DEBUG
-log_format = "%(log_color)s%(levelname)-8s%(reset)s  %(log_color)s%(message)s%(reset)s"
-
-logging.root.setLevel(log_level)
-formatter = ColoredFormatter(log_format)
-
-stream = logging.StreamHandler()
-stream.setLevel(log_level)
-stream.setFormatter(formatter)
-
-LOG = logging.getLogger("pythonConfig")
-LOG.setLevel(log_level)
-LOG.addHandler(stream)
-
-
-# If the app is running in Docker
-if "BASE_URL" in os.environ or "IP_TO_LOCATION" in os.environ:
-    BASE_URL = os.environ["BASE_URL"]
-    IP_TO_LOCATION = (
-        True if os.environ["IP_TO_LOCATION"].upper() in ["TRUE", "T"] else False
-    )
-    if IP_TO_LOCATION:
-        API_KEY = os.environ["API_KEY"]
-    else:
-        API_KEY = None
-
-# Otherwise, the app is running on bare metal
-try:
-    with open("config.ini", "r") as f:
-        config = configparser.ConfigParser()
-        config.read_string(f.read())
-
-        BASE_URL = config["CONFIG"]["BASE_URL"]
-        IP_TO_LOCATION = (
-            True
-            if config["CONFIG"]["IP_TO_LOCATION"].upper() in ["TRUE", "T"]
-            else False
-        )
-        if IP_TO_LOCATION:
-            API_KEY = config["CONFIG"]["API_KEY"]
-        else:
-            API_KEY = None
-except FileNotFoundError:
-    config = configparser.ConfigParser()
-    config["CONFIG"] = {"BASE_URL": "", "IP_TO_LOCATION": "", "API_KEY": ""}
-
-    with open("config.ini", "w") as configfile:
-        config.write(configfile)
-
-    LOG.error(
-        "`config.ini` has been created. Fill out the necessary information then re-run."
-    )
-    exit()