Add args to ipapi
This commit is contained in:
parent
0a20cae93f
commit
c12f18cd2c
33
ipapi.py
33
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
|
||||
uvicorn.run(app, host="0.0.0.0", port=8000)
|
||||
main()
|
||||
|
|
Loading…
Reference in New Issue