From 026618b8d14ad4f2586e95a8429dbd6a80360ab0 Mon Sep 17 00:00:00 2001
From: Scott Wallace <scott@wallace.sh>
Date: Tue, 28 Dec 2021 15:15:08 +0000
Subject: [PATCH] Add IP-based protection

---
 main.py       |  3 ++-
 slinky/web.py | 27 ++++++++++++++++++++++++++-
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/main.py b/main.py
index 2b6e1e4..1dc132f 100644
--- a/main.py
+++ b/main.py
@@ -4,7 +4,7 @@ Main Flask-based app for Slinky
 from flask import Flask, render_template
 from flask_bootstrap import Bootstrap
 
-from slinky.web import slinky_webapp
+from slinky.web import protect, slinky_webapp
 
 app = Flask(__name__)
 app.register_blueprint(slinky_webapp)
@@ -13,6 +13,7 @@ Bootstrap(app)
 
 
 @app.route('/')
+@protect
 def index() -> str:
     """
     Index/Landing page
diff --git a/slinky/web.py b/slinky/web.py
index 755a459..98342c8 100644
--- a/slinky/web.py
+++ b/slinky/web.py
@@ -4,9 +4,11 @@ Web component
 
 import logging
 from datetime import datetime
+from functools import wraps
+from typing import Any, Callable
 
 import yaml
-from flask import Blueprint, Response, redirect, render_template
+from flask import Blueprint, Response, redirect, render_template, request
 from flask_wtf import FlaskForm
 from wtforms import DateTimeLocalField, HiddenField, IntegerField, StringField
 from wtforms.validators import DataRequired, Length
@@ -67,6 +69,27 @@ class AddForm(FlaskForm):  # type: ignore[misc]
     )
 
 
+def protect(func: Callable[..., Response]) -> Callable[..., Response]:
+    """
+    Decorator that will protect the admin interface
+
+    Args:
+        func (Callable): Wrapped function
+
+    Returns:
+        Callable: Function wrapper
+    """
+
+    @wraps(func)
+    def check_ip(*args: Any, **kwargs: Any) -> Response:
+        if request.remote_addr not in cfg['allowed_ips']:
+            print(f'Protected URL access attempt from {request.remote_addr}')
+            return Response('Not found', 404)
+        return func(*args, **kwargs)
+
+    return check_ip
+
+
 @slinky_webapp.route('/<path:path>')
 def try_path_as_shortcode(path: str) -> Response:
     """
@@ -95,6 +118,7 @@ def try_path_as_shortcode(path: str) -> Response:
 
 
 @slinky_webapp.route('/_/add', methods=['GET', 'POST'])
+@protect
 def add() -> Response:
     """
     Create and add a new shorturl
@@ -130,6 +154,7 @@ def add() -> Response:
 
 
 @slinky_webapp.route('/_/list', methods=['GET', 'POST'])
+@protect
 def lister() -> str:
     """
     Create and add a new shorturl