more error checking, fixed a todo

.gitignore

@@ -0,0 +1,4 @@
+.vscode
+__pycache__
+app/mirrors.yaml
+app/userdata.csv

app/main.py

@@ -1,140 +1,131 @@
 """ flask app to redirect request to appropriate armbian mirror and image """
 
 import json
-import uwsgi
 
+#import uwsgi
 from flask import (
         Flask,
         redirect,
-        request
+        request,
+        Response,
 )
-# from markupsafe import escape
+from geolite2 import geolite2
 from download_image_map import Parser
 from mirror_list import Mirror
-from geolite2 import geolite2
-from ruamel.yaml import YAML
-# from ruamel.yaml.scalarstring import (
-#     FoldedScalarString,
-#     LiteralScalarString,
-# )
 
 
-def load_mirrors():
-    """ open mirrors file and return contents """
-    global mode
-    yaml = YAML()
-    yaml.indent(mapping=2, sequence=4, offset=2)
-    yaml.preserve_quotes = True
+mirror = Mirror()
+if mirror.mode == "dl_map":
+    parser = Parser('userdata.csv')
+    DL_MAP = parser.parsed_data
+else:
+    DL_MAP = None
 
-    with open('mirrors.yaml', 'r') as f:
-        config = yaml.load(f)
-    mode = config['mode']
-    print("using mode: {}".format(mode))
-    return config['mirrors']
-
-
-def reload_all():
-    """ reload mirror and redirect map files """
-    global mode
-    mirror = Mirror(load_mirrors())
-    if mode == "dl_map":
-        global dl_map
-        dl_map = parser.reload()
-    return mirror
+geolite_reader = geolite2.reader()
+app = Flask(__name__)
 
+# def reload_all():
+#     """ reload mirror and redirect map files """
+#     mirror = Mirror()
+#     if mirror.mode == "dl_map":
+#         global dl_map
+#         dl_map = parser.reload()
+#     return mirror
 
 def get_ip():
-    """ returns requestor's IP by parsersing proxy  headers """
-    if request.environ.get('HTTP_X_FORWARDED_FOR') is None:
-        return request.environ['REMOTE_ADDR']
-    return request.environ['HTTP_X_FORWARDED_FOR']
+    """ returns client IP by parsing proxy headers
+        if they don't exist, return the actual client IP address """
+    return request.environ.get('HTTP_X_FORWARDED_FOR',
+                               request.environ.get('REMOTE_ADDR'),
+                              )
 
 
-def get_region(IP):
+def get_region(client_ip, reader=geolite_reader, continents=mirror.continents):
     """ this is where we geoip and return region code """
+    if client_ip.startswith(("192.168.", "10.")):
+        print(f"Local IP address: {client_ip}")
+        return None
     try:
-        match = reader.get(IP)
-        conti = match['continent']['code']
-# FIXME Get Contient List from Configuration File
-        if conti in ("EU", "NA", "AS"):
-            print("Match {} to continent {}".format(IP, conti))
+        match = reader.get(client_ip)
+        if not match:
+            print(f"match failure for IP: {client_ip}")
+            return None
+        # matches are a dict where we want match["continent"]["code"] = "EU"
+        conti = match.get("continent", {}).get("code")
+        if conti in continents:
+            print(f"Match {client_ip} to continent {conti}")
             return conti
-    except:
-        print("match failure for IP: {}".format(IP))
+    # pylint: disable=broad-except
+    except Exception as error_message:
+        print(f"match failure for IP: {client_ip} (Error: {error_message}")
         print(json.dumps(match))
     else:
         return None
 
-
-def get_redirect(path, IP):
+def get_redirect(path, client_ip, mirror_class=mirror, dl_map=DL_MAP):
     """ get redirect based on path and IP """
-    global mode
-    global dl_map
-    region = get_region(IP)
+    region = get_region(client_ip)
     split_path = path.split('/')
+
     if split_path[0] == "region":
-        if split_path[1] in mirror.all_regions():
+        if split_path[1] in mirror_class.all_regions():
             region = split_path[1]
         del split_path[0:2]
         path = "{}".format("/".join(split_path))
-    if mode == "dl_map" and len(split_path) == 2:
+
+    if mirror_class.mode == "dl_map" and len(split_path) == 2:
         key = "{}/{}".format(split_path[0], split_path[1])
         new_path = dl_map.get(key, path)
-        return "{}{}".format(mirror.next(region), new_path)
+        return "{}{}".format(mirror_class.next(region), new_path)
+
     if path == '':
-        return mirror.next(region)
-    return "{}{}".format(mirror.next(region), path)
+        return mirror_class.next(region)
+
+    return "{}{}".format(mirror_class.next(region), path)
 
 
-mirror = Mirror(load_mirrors())
-if mode == "dl_map":
-    parser = Parser('userdata.csv')
-    dl_map = parser.parsed_data
-
-reader = geolite2.reader()
-app = Flask(__name__)
-
-
-@ app.route('/status')
+@app.route('/status')
 def status():
     """ return health check status """
-    return "OK"
+    resp = Response("OK")
+    resp.headers['X-Client-IP'] = get_ip()
+    return resp
+
+# @app.route('/reload')
+# def signal_reload():
+#     """ trigger graceful reload via uWSGI """
+#     uwsgi.reload()
+#     return "reloaded"
 
 
-@ app.route('/reload')
-def signal_reload():
-    """ trigger graceful reload via uWSGI """
-    uwsgi.reload()
-    return "reloding"
-
-
-@ app.route('/mirrors')
+@app.route('/mirrors')
 def show_mirrors():
     """ return all_mirrors in json format to requestor """
     return json.dumps(mirror.all_mirrors())
 
 
-@ app.route('/regions')
+@app.route('/regions')
 def show_regions():
     """ return all_regions in json format to requestor """
     return json.dumps(mirror.all_regions())
 
 
-@ app.route('/dl_map')
-def show_dl_map():
-    global mode
-    global dl_map
-    if mode == "dl_map":
+@app.route('/dl_map')
+def show_dl_map(mirror_mode=mirror.mode, dl_map=DL_MAP):
+    """ returns a direct-download map """
+    if mirror_mode == "dl_map":
         return json.dumps(dl_map)
     return "no map. in direct mode"
 
 
-@ app.route('/geoip')
-def show_geoip():
+@app.route('/geoip')
+def show_geoip(reader=geolite_reader):
+    """ returns the geoip location of the client IP """
     return json.dumps(reader.get(get_ip()))
 
-@ app.route('/', defaults={'path': ''})
-@ app.route('/<path:path>')
+
+@app.route('/', defaults={'path': ''})
+@app.route('/<path:path>')
 def catch_all(path):
     """ default app route for redirect """
     return redirect(get_redirect(path, get_ip()), 302)

app/mirror_list.py

@@ -1,23 +1,39 @@
 """ manage iterating through mirrors and return appropriate one based by region """
 
+from ruamel.yaml import YAML
 
 class Mirror():
 
-    def __init__(self, mirror_list):
-        self.mirror_list = mirror_list
+    def __init__(self):
+        self.load_mirrors()
         self._list_position = dict()
         self._list_max = dict()
+        self.continents = [region for region in self.mirror_list.keys() ]
+        self.regions = self.continents + [ 'default' ]
         self.mirror_list['default'] = list(
-            mirror_list['NA'] + mirror_list['EU'])
-        for region in list(mirror_list.keys()):
+            self.mirror_list['NA'] + self.mirror_list['EU'])
+        for region in self.regions:
             self._list_position[region] = 0
-            self._list_max[region] = len(mirror_list[region]) - 1
+            self._list_max[region] = len(self.mirror_list[region]) - 1
+
+    def load_mirrors(self):
+        """ open mirrors file and return contents """
+        yaml = YAML()
+        yaml.indent(mapping=2, sequence=4, offset=2)
+        yaml.preserve_quotes = True
+
+        with open('mirrors.yaml', 'r') as f:
+            config = yaml.load(f)
+        self.mode = config['mode']
+        print("using mode: {}".format(self.mode))
+        self.mirror_list = config.get('mirrors', {})
+        return self.mirror_list
 
     # set defaults to None in param list, then actually set inside
     # body to avoid scope change
     def increment(self, region=None):
         """ move to next element regions mirror list and return list position"""
-        if region is None:
+        if region is None or region not in self.all_regions():
             region = 'default'
         if self._list_position[region] == self._list_max[region]:
             self._list_position[region] = 0
@@ -29,7 +45,7 @@ class Mirror():
         if region is None:
             region = 'default'
         self.increment(region)
-        return self.mirror_list[region][self._list_position[region]]
+        return self.mirror_list.get(region, self.mirror_list.get('default'))[self._list_position[region]]
 
     def all_mirrors(self):
         """ return all mirrrors configured """
@@ -37,4 +53,4 @@ class Mirror():
 
     def all_regions(self):
         """ return list of regions configured """
-        return list(self.mirror_list.keys())
+        return self.regions

app/requirements.txt

@@ -1,8 +1,6 @@
 flask
 uwsgi
-#boto3
 python-dotenv
-uwsgi
 ruamel.yaml
 maxminddb
 maxminddb-geolite2

app/test_flask_app.py

@@ -0,0 +1,66 @@
+import os
+import tempfile
+import json
+
+import pytest
+
+from main import app
+
+
+TEST_IPS = [ 
+        "10.0.0.1",
+        "1.2.3.4",
+        "123.254.123.233",
+        "5.254.123.233",
+    ]
+
+@pytest.fixture
+def client():
+    with app.test_client() as client:
+        yield client
+
+def test_status(client):
+    """ test the status endpoint """
+
+    rv = client.get('/status')
+    assert rv.data == b'OK'
+
+def test_mirrors(client):
+    """ test the mirrors endpoint """
+
+    rv = client.get('/mirrors')
+    jsondata = rv.data.decode('utf-8')
+
+    assert isinstance(json.loads(jsondata), dict)
+
+def test_regions(client):
+    """ test the regions endpoint """
+
+    rv = client.get('/regions')
+    jsondata = rv.data.decode('utf-8')
+
+    assert isinstance(json.loads(jsondata), list)
+
+def test_regions_proxied(client):
+    """ test the regions endpoint, but send x-forwarded-for headers """
+
+    for address in TEST_IPS:
+        test_headers = {
+            "X-Forwarded-For" : address,
+            }
+        rv = client.get('/regions', headers=test_headers)
+        jsondata = rv.data.decode('utf-8')
+        assert isinstance(json.loads(jsondata), list)
+
+def test_status_proxied(client):
+    """ test the status endpoint, but send x-forwarded-for headers 
+        the status endpoint sends an X-Client-IP header in the response
+        """
+
+    for address in TEST_IPS:
+        test_headers = {
+            "X-Forwarded-For" : address,
+            }
+        rv = client.get('/status', headers=test_headers)
+        assert 'X-Client-IP' in rv.headers
+        assert rv.headers.get('X-Client-IP') == address

run_local_dev.sh

@@ -1,15 +1,34 @@
-#export FLASK_APP=main.py
-#python -m flask run --host 0.0.0.0 --port 5000
-
-
+#!/bin/bash
 
 APP_PATH=$(pwd)/app
-USERDATA_PATH=$(pwd)/userdata.csv
-MIRRORS_CONF_PATH=$(pwd)/mirrors.yaml
+USERDATA_PATH=$(pwd)/examples/userdata.csv
+MIRRORS_CONF_PATH=$(pwd)/examples/mirrors-apt.yaml
 LISTEN_PORT=5000
 CONTAINER_NAME=redirect
 DETACH=false
 
 ##FIXME CHANGE CONFIG MAP TO YAML WHEN DONE
 
-sudo docker run --rm $([[ ${DETACH} == "true" ]] && echo "-d") -v ${APP_PATH}:/app -v ${USERDATA_PATH}:/app/userdata.csv -v ${MIRRORS_CONF_PATH}:/app/mirrors.yaml -p ${LISTEN_PORT}:80 --name ${CONTAINER_NAME} quay.io/lanefu/nginx-uwsgi-flask:arm64
+if [ ! -d "${APP_PATH}" ]; then
+    echo "Unable to find App path: ${APP_PATH}"
+    exit 1
+fi
+
+if [ ! -f "${USERDATA_PATH}" ]; then
+    echo "Unable to find userdata.csv at ${USERDATA_PATH}"
+    exit 1
+fi
+
+if [ ! -f "${MIRRORS_CONF_PATH}" ]; then
+    echo "Unable to find mirrors.yaml at ${MIRRORS_CONF_PATH}"
+    exit 1
+fi
+
+
+sudo docker run --rm $([[ ${DETACH} == "true" ]] && echo "-d") \
+    -v ${APP_PATH}:/app \
+    -v ${USERDATA_PATH}:/app/userdata.csv \
+    -v ${MIRRORS_CONF_PATH}:/app/mirrors.yaml \
+    -p ${LISTEN_PORT}:80 \
+    --name ${CONTAINER_NAME} \
+    quay.io/lanefu/nginx-uwsgi-flask:arm64

run_tests.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+
+APP_PATH=$(pwd)/app
+USERDATA_PATH=$(pwd)/examples/userdata.csv
+MIRRORS_CONF_PATH=$(pwd)/examples/mirrors-apt.yaml
+LISTEN_PORT=5000
+CONTAINER_NAME=redirect_test
+DETACH=false
+
+##FIXME CHANGE CONFIG MAP TO YAML WHEN DONE
+
+if [ ! -d "${APP_PATH}" ]; then
+    echo "Unable to find App path: ${APP_PATH}"
+    exit 1
+fi
+
+if [ ! -f "${USERDATA_PATH}" ]; then
+    echo "Unable to find userdata.csv at ${USERDATA_PATH}"
+    exit 1
+fi
+
+if [ ! -f "${MIRRORS_CONF_PATH}" ]; then
+    echo "Unable to find mirrors.yaml at ${MIRRORS_CONF_PATH}"
+    exit 1
+fi
+
+sudo docker run --rm $([[ ${DETACH} == "true" ]] && echo "-d") \
+    -v ${APP_PATH}:/app \
+    -v ${USERDATA_PATH}:/app/userdata.csv \
+    -v ${MIRRORS_CONF_PATH}:/app/mirrors.yaml \
+    -p ${LISTEN_PORT}:80 \
+    --name ${CONTAINER_NAME} \
+    quay.io/lanefu/nginx-uwsgi-flask:arm64 bash -c "pip install --upgrade pip && pip install -r requirements.txt && pip install pytest && pytest -s "
+