Add args to ipapi

ipapi.py

@@ -1,12 +1,15 @@
 #!/usr/bin/env python3
 
+import re
+import subprocess
+import argparse
+from dataclasses import dataclass
+import uvicorn
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.middleware.cors import CORSMiddleware
 from slowapi.errors import RateLimitExceeded
 from slowapi import Limiter, _rate_limit_exceeded_handler
 from slowapi.util import get_remote_address
-import subprocess
-import re
 
 limiter = Limiter(key_func=get_remote_address)
 app = FastAPI()
@@ -22,8 +25,10 @@ app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
 
 ip_regex = re.compile(r"^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$")
 
+limit = "1/second"
+
 @app.get("/api/rdns/{ip}")
-@limiter.limit("1/second")
+@limiter.limit(lambda: limit)
 async def get_rdns(ip: str, request: Request):
     match = re.match(ip_regex, ip)
     if match is None:
@@ -34,6 +39,24 @@ async def get_rdns(ip: str, request: Request):
         raise HTTPException(status_code=404, detail="No rDNS information found for IP")
     return { "ip": sanitized_ip, "rdns": output.rstrip(".") }
 
+@dataclass
+class IpApiArgs:
+    address: str
+    port: int
+    limit: str
+
+def parse_list_arg(arg: str):
+    return [x.strip().lower() for x in arg.split(",") if x.strip()]
+
+def main():
+    global limit
+    parser = argparse.ArgumentParser("ipmap")
+    parser.add_argument("-a", "--address", default = "127.0.0.1", help = "the address to use for the api (default: %(default)s)")
+    parser.add_argument("-p", "--port", default = 8000, type = int, help = "the port to use for the api (default: %(default)s)")
+    parser.add_argument("-l", "--limit", default = limit, help = "the rate limit for the api (default: %(default)s)")
+    args = parser.parse_args(namespace = IpApiArgs)
+    limit = args.limit
+    uvicorn.run(app, host=args.address, port=args.port)
+
 if __name__ == "__main__":
-    import uvicorn
+    main()
-    uvicorn.run(app, host="0.0.0.0", port=8000)