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('/') 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