Block undesired IP addresses

rpipe/__init__.py

@@ -1 +1 @@
-__version__: str = "9.5.4"  # Must be "<major>.<minor>.<patch>", all numbers
+__version__: str = "9.6.0"  # Must be "<major>.<minor>.<patch>", all numbers

rpipe/client/admin.py

@@ -145,6 +145,21 @@ def unlock(self) -> None:
         """
         self._lock(False)
 
+    def ip(self, block: str | None, unblock: str | None) -> None:
+        """
+        Request the blocked ip addresses, or block / unblock an ip address
+        """
+        if block is not None and unblock is not None:
+            raise ValueError("block and unblock may not both be non-None")
+        if block is None and unblock is None:
+            blocked = self._request("/admin/ip", '{"ip": null}').text
+            print(f"Blocked IP addresses: {blocked}")
+            return
+        ban = block is not None
+        addr = block if ban else unblock
+        self._request("/admin/ip", dumps({"ip": addr, "block": ban}))
+        print(f"{"" if ban else "UN"}BLOCKED: {addr}")
+
 
 class Admin:
     """

rpipe/client/cli.py

@@ -177,5 +177,9 @@ def cli() -> None:
     log_lvl_p.add_argument("level", default=None, nargs="?", help="The log level for the server to use")
     admin.add_parser("lock", help="Lock the channel")
     admin.add_parser("unlock", help="Unlock the channel")
+    ip_p = admin.add_parser("ip", help="Block / unblock ip addresses, or get a list of blocked addresses")
+    m_g = ip_p.add_mutually_exclusive_group(required=False)
+    m_g.add_argument("--block", help="Block a given IP address")
+    m_g.add_argument("--unblock", help="Unblock a given IP address")
     argcomplete.autocomplete(parser)  # Tab completion
     _cli(parser, parser.parse_args())

rpipe/server/admin/admin.py

@@ -13,6 +13,7 @@
 
 if TYPE_CHECKING:
     from pathlib import Path
+    from ..app import Blocked
     from ..server import State
 
 
@@ -24,9 +25,10 @@ class Methods:
     Protected methods that can be accessed by the Admin class
     """
 
-    def __init__(self, log_file: Path | None) -> None:
+    def __init__(self, log_file: Path | None, blocked: Blocked) -> None:
         self._log = getLogger("Admin")
         self._log_file = log_file
+        self._blocked = blocked
         self._log.debug("Log file set to %s", log_file)
 
     @staticmethod
@@ -63,7 +65,8 @@ def stats(state: State, _: str) -> Response:
             stats = asdict(s.stats)
         return json_response(stats)
 
-    def channels(self, state: State, _: str) -> Response:
+    @staticmethod
+    def channels(state: State, _: str) -> Response:
         """
         Return a list of the server's current channels and stats
         """
@@ -81,6 +84,21 @@ def lock(self, state: State, body: str) -> Response:
             s.locked = lock
         return Response(f"Channel {channel} is now {lock_s}", status=200)
 
+    def ip(self, _: State, body: str) -> Response:
+        js = loads(body.strip())
+        if (addr := js["ip"]) is None:
+            return json_response(self._blocked.data["ips"])
+        lst = self._blocked.data["ips"]
+        if js["block"]:
+            if addr not in lst:
+                lst.append(addr)
+                self._blocked.commit()
+        elif addr in lst:
+            while addr in lst:
+                lst.remove(addr)
+            self._blocked.commit()
+        return Response(status=200)
+
 
 class Admin:
     """
@@ -92,9 +110,9 @@ class Admin:
 
     __slots__ = ("_verify", "_methods")
 
-    def __init__(self, log_file: Path, key_files: list[Path]) -> None:
+    def __init__(self, log_file: Path, key_files: list[Path], blocked: Blocked) -> None:
         self._verify = Verify(key_files)
-        self._methods = Methods(log_file)
+        self._methods = Methods(log_file, blocked)
 
     def __getattr__(self, item: str) -> Any:
         """

rpipe/server/app.py

@@ -6,6 +6,8 @@
 from functools import wraps
 from pathlib import Path
 import atexit
+import typing
+import json
 
 from flask import Response, Flask, send_file, request
 from zstdlib.log import CuteFormatter
@@ -35,15 +37,43 @@ class ServerConfig:
     port: int
     debug: bool
     state_file: Path | None
+    block_file: Path | None
     key_files: list[Path]
 
 
+class Blocked:
+    _DEFAULT: dict[str, list[str]] = {"ips": [], "routes": []}
+
+    def __init__(self, file: Path | None) -> None:
+        self.data = dict(self._DEFAULT) if file is None else json.loads(file.read_text())
+        self.file: Path | None = file
+        self._lg = getLogger("Blocked")
+
+    def commit(self) -> None:
+        if self.file is None:
+            raise ValueError("Cannot save a block file when block-file not set")
+        self.file.write_text(json.dumps(self.data))
+
+    def __call__(self) -> bool:
+        if self.file is None:
+            return False
+        if (ip := request.remote_addr) in self.data["ips"]:
+            return True
+        if (pth := request.path) in self.data["routes"]:
+            self._lg.info("Blocking IP %s based on route: %s", ip, pth)
+            self.data["ips"].append(typing.cast(str, ip))
+            self.commit()
+            return True
+        return False
+
+
 class App(Flask):
 
     @dataclass(frozen=True, slots=True)
     class Objs:
         admin: Admin
         server: Server
+        blocked: Blocked
         favicon: Path | None
 
     def __init__(self) -> None:
@@ -57,10 +87,11 @@ def start(self, conf: ServerConfig, log_file: Path, favicon: Path | None):
         if favicon is not None and not favicon.is_file():
             lg.error("Favicon file not found: %s", favicon)
             favicon = None
-        admin = Admin(log_file, conf.key_files)
+        blocked = Blocked(conf.block_file)
+        admin = Admin(log_file, conf.key_files, blocked)
         lg.info("Starting server version: %s", __version__)
         # pylint: disable=attribute-defined-outside-init
-        self._objs = self.Objs(admin, Server(conf.debug, conf.state_file), favicon)
+        self._objs = self.Objs(admin, Server(conf.debug, conf.state_file), blocked, favicon)
         lg.info("Binding to %s:%s", conf.host, conf.port)
         if conf.debug:
             self.run(host=conf.host, port=conf.port, debug=True)
@@ -70,12 +101,15 @@ def start(self, conf: ServerConfig, log_file: Path, favicon: Path | None):
     def give(self, *, objs: bool = False, logged: bool = True):
         """
         Give the wrapped function self.objs and log requests as requested
+        Also handles blocked IP addresses and routes
         """
         lg = getLogger(_LOG)
 
         def decorator(func):
             @wraps(func)
             def inner(*args, **kwargs):
+                if self._objs.blocked():
+                    return Response(status=401)
                 ret = func(*args, self._objs, **kwargs) if objs else func(*args, **kwargs)
                 if not logged or self._objs.server.state.debug:
                     return ret
@@ -214,6 +248,11 @@ def _admin_lock(o: App.Objs) -> Response:
     return o.admin.lock(o.server.state)
 
 
+@app.route("/admin/ip", admin=True)
+def _admin_ip(o: App.Objs) -> Response:
+    return o.admin.ip(o.server.state)
+
+
 # Main functions
 
 

rpipe/server/main.py

@@ -19,6 +19,7 @@ def cli() -> None:
     )
     parser.add_argument("port", type=int, help="The port waitress will listen on")
     parser.add_argument("--host", default="0.0.0.0", help="The host waitress will bind to for listening")
+    parser.add_argument("--block-file", type=Path, help="A json of IP addresses and routes to ban")
     parser.add_argument("-s", "--state-file", type=Path, help="The save state file, if desired")
     parser.add_argument(
         "-k",