Replace dig command with dnspython library

This commit is contained in:
LilyRose2798 2024-04-11 17:45:56 +10:00
parent 5fb2db4ab1
commit 8675ce6c08
3 changed files with 40 additions and 10 deletions

View File

@ -1,11 +1,14 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
import re import re
import subprocess
import argparse import argparse
from dataclasses import dataclass from dataclasses import dataclass
import dns.exception
import dns.rdatatype
import dns.resolver
import dns.reversename
import uvicorn import uvicorn
from fastapi import FastAPI, HTTPException, Request from fastapi import FastAPI, HTTPException, Request, Response
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
from slowapi.errors import RateLimitExceeded from slowapi.errors import RateLimitExceeded
from slowapi import Limiter, _rate_limit_exceeded_handler from slowapi import Limiter, _rate_limit_exceeded_handler
@ -29,15 +32,21 @@ limit = "1/second"
@app.get("/api/rdns/{ip}") @app.get("/api/rdns/{ip}")
@limiter.limit(lambda: limit) @limiter.limit(lambda: limit)
async def get_rdns(ip: str, request: Request): async def get_rdns(ip: str, request: Request, response: Response):
match = re.match(ip_regex, ip) try:
if match is None: answer = dns.resolver.resolve_address(ip, search = True)
except dns.exception.SyntaxError:
raise HTTPException(status_code=400, detail="Invalid IP address") raise HTTPException(status_code=400, detail="Invalid IP address")
sanitized_ip = match[0] except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
output = subprocess.run(["dig", "-x", sanitized_ip, "+short"], capture_output = True, text = True).stdout.strip()
if not output:
raise HTTPException(status_code=404, detail="No rDNS information found for IP") raise HTTPException(status_code=404, detail="No rDNS information found for IP")
return { "ip": sanitized_ip, "rdns": output.rstrip(".") } except dns.resolver.LifetimeTimeout:
raise HTTPException(status_code=504, detail="Request for rDNS information timed out")
except dns.resolver.NoNameservers:
raise HTTPException(status_code=503, detail="No nameservers currently available to fulfil request")
except:
raise HTTPException(status_code=500, detail="Unexpected error occurred")
response.headers["Cache-Control"] = f"max-age={answer.rrset.ttl}"
return { "ip": ip, "rdns": str(answer[0]).rstrip(".") }
@dataclass @dataclass
class IpApiArgs: class IpApiArgs:

22
poetry.lock generated
View File

@ -93,6 +93,26 @@ wrapt = ">=1.10,<2"
[package.extras] [package.extras]
dev = ["PyTest", "PyTest-Cov", "bump2version (<1)", "sphinx (<2)", "tox"] dev = ["PyTest", "PyTest-Cov", "bump2version (<1)", "sphinx (<2)", "tox"]
[[package]]
name = "dnspython"
version = "2.6.1"
description = "DNS toolkit"
optional = false
python-versions = ">=3.8"
files = [
{file = "dnspython-2.6.1-py3-none-any.whl", hash = "sha256:5ef3b9680161f6fa89daf8ad451b5f1a33b18ae8a1c6778cdf4b43f08c0a6e50"},
{file = "dnspython-2.6.1.tar.gz", hash = "sha256:e8f0f9c23a7b7cb99ded64e6c3a6f3e701d78f50c55e002b839dea7225cff7cc"},
]
[package.extras]
dev = ["black (>=23.1.0)", "coverage (>=7.0)", "flake8 (>=7)", "mypy (>=1.8)", "pylint (>=3)", "pytest (>=7.4)", "pytest-cov (>=4.1.0)", "sphinx (>=7.2.0)", "twine (>=4.0.0)", "wheel (>=0.42.0)"]
dnssec = ["cryptography (>=41)"]
doh = ["h2 (>=4.1.0)", "httpcore (>=1.0.0)", "httpx (>=0.26.0)"]
doq = ["aioquic (>=0.9.25)"]
idna = ["idna (>=3.6)"]
trio = ["trio (>=0.23)"]
wmi = ["wmi (>=1.5.1)"]
[[package]] [[package]]
name = "fastapi" name = "fastapi"
version = "0.110.1" version = "0.110.1"
@ -560,4 +580,4 @@ files = [
[metadata] [metadata]
lock-version = "2.0" lock-version = "2.0"
python-versions = "^3.11" python-versions = "^3.11"
content-hash = "2afa2116759771dde6755f47cffb03bec0527a75df805fa9a9fd8cc2a9d8fde0" content-hash = "882810214ec005c8e1d0b99099d0f9fc8d6e8fb9140ac9f452e18e7e3c580176"

View File

@ -15,6 +15,7 @@ cmap = "^0.1.3"
fastapi = "^0.110.1" fastapi = "^0.110.1"
uvicorn = "^0.29.0" uvicorn = "^0.29.0"
slowapi = "^0.1.9" slowapi = "^0.1.9"
dnspython = "^2.6.1"
[build-system] [build-system]