diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ed7d1b6 --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +.vscode +__pycache__ +app/mirrors.yaml +app/userdata.csv \ No newline at end of file diff --git a/app/main.py b/app/main.py index 584a8f7..57b1eb4 100644 --- a/app/main.py +++ b/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('/') + +@app.route('/', defaults={'path': ''}) +@app.route('/') def catch_all(path): """ default app route for redirect """ return redirect(get_redirect(path, get_ip()), 302) diff --git a/app/mirror_list.py b/app/mirror_list.py index c796657..9886e45 100644 --- a/app/mirror_list.py +++ b/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 diff --git a/app/requirements.txt b/app/requirements.txt index 96820a8..7d3d4d4 100644 --- a/app/requirements.txt +++ b/app/requirements.txt @@ -1,8 +1,6 @@ flask uwsgi -#boto3 python-dotenv -uwsgi ruamel.yaml maxminddb maxminddb-geolite2 diff --git a/app/test_flask_app.py b/app/test_flask_app.py new file mode 100755 index 0000000..2dcdd98 --- /dev/null +++ b/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 diff --git a/run_local_dev.sh b/run_local_dev.sh index b21b5f8..7d176a7 100755 --- a/run_local_dev.sh +++ b/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 diff --git a/run_tests.sh b/run_tests.sh new file mode 100755 index 0000000..41324a9 --- /dev/null +++ b/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 " +