Replace dig command with dnspython library

ipapi.py

@@ -1,11 +1,14 @@
 #!/usr/bin/env python3
 
 import re
-import subprocess
 import argparse
 from dataclasses import dataclass
+import dns.exception
+import dns.rdatatype
+import dns.resolver
+import dns.reversename
 import uvicorn
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException, Request, Response
 from fastapi.middleware.cors import CORSMiddleware
 from slowapi.errors import RateLimitExceeded
 from slowapi import Limiter, _rate_limit_exceeded_handler
@@ -29,15 +32,21 @@ limit = "1/second"
 
 @app.get("/api/rdns/{ip}")
 @limiter.limit(lambda: limit)
-async def get_rdns(ip: str, request: Request):
-    match = re.match(ip_regex, ip)
-    if match is None:
+async def get_rdns(ip: str, request: Request, response: Response):
+    try:
+        answer = dns.resolver.resolve_address(ip, search = True)
+    except dns.exception.SyntaxError:
         raise HTTPException(status_code=400, detail="Invalid IP address")
-    sanitized_ip = match[0]
-    output = subprocess.run(["dig", "-x", sanitized_ip, "+short"], capture_output = True, text = True).stdout.strip()
-    if not output:
+    except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
         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
 class IpApiArgs:

poetry.lock

@@ -93,6 +93,26 @@ wrapt = ">=1.10,<2"
 [package.extras]
 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]]
 name = "fastapi"
 version = "0.110.1"
@@ -560,4 +580,4 @@ files = [
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "2afa2116759771dde6755f47cffb03bec0527a75df805fa9a9fd8cc2a9d8fde0"
+content-hash = "882810214ec005c8e1d0b99099d0f9fc8d6e8fb9140ac9f452e18e7e3c580176"

pyproject.toml

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