Add args to ipapi

This commit is contained in:
LilyRose2798 2024-04-11 16:31:47 +10:00
parent 0a20cae93f
commit c12f18cd2c
1 changed files with 28 additions and 5 deletions

View File

@ -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
uvicorn.run(app, host="0.0.0.0", port=8000)
main()