From 606e700f11b2b002de8410d75935c7408087ea21 Mon Sep 17 00:00:00 2001
From: nexy7574 <git@nexy7574.co.uk>
Date: Sun, 28 Apr 2024 11:43:14 +0100
Subject: [PATCH] fixes

---
 ipserv.py | 76 ++++++++++++++++++++++++++++++-------------------------
 1 file changed, 41 insertions(+), 35 deletions(-)

diff --git a/ipserv.py b/ipserv.py
index 7178571..63c9a72 100644
--- a/ipserv.py
+++ b/ipserv.py
@@ -36,7 +36,8 @@ async def lifespan(_app: FastAPI):
     async with aiohttp.ClientSession(
         "https://ip.shronk.tech",
         connector=aiohttp.TCPConnector(limit=2048),
-        raise_for_status=True
+        raise_for_status=True,
+        timeout=aiohttp.ClientTimeout(5)
     ) as client:
         _app.state.session = client
         yield
@@ -85,6 +86,7 @@ async def make_request(ip: str, headers: dict[str, str] = None) -> dict | HTTPEx
         logging.info("cache expired for %s", ip)
 
     try:
+        logging.info("Querying SHRoNKNet: %s", ip)
         async with app.state.session.get(f"/lookup?ip={ip}", headers=headers) as response:
             data = await response.json()
             data["source"] = "SHRoNKNet"
@@ -99,23 +101,25 @@ async def make_request(ip: str, headers: dict[str, str] = None) -> dict | HTTPEx
 
 async def _get_ipinfo_no_token(ip: str):
     try:
-        async with app.state.session.get("https://ipinfo.io/widget/demo/" + ip) as response:
-            data = (await response.json())["data"]
-            return {
-                "ip": data["ip"],
-                "city": data["city"],
-                "country": data["region"],
-                "countryCode": data["country"],
-                "asn": int(data["asn"]["asn"][2:]),
-                "isp": data["asn"]["name"],
-                "source": "ipinfo",
-                "lat": data["loc"].split(",")[0],
-                "lon": data["loc"].split(",")[1],
-                "hostname": data["hostname"],
-                "timezone": data["timezone"],
-                "subnet": data["abuse"]["network"],
-                "abuse": data["abuse"]["email"]
-            }
+        logging.info("Querying IPInfo WITHOUT TOKEN: %s", ip)
+        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(5)) as session:
+            async with session.get("https://ipinfo.io/widget/demo/" + ip) as response:
+                data = (await response.json())["data"]
+                return {
+                    "ip": data["ip"],
+                    "city": data["city"],
+                    "country": data["region"],
+                    "countryCode": data["country"],
+                    "asn": int(data["asn"]["asn"][2:]),
+                    "isp": data["asn"]["name"],
+                    "source": "ipinfo",
+                    "lat": data["loc"].split(",")[0],
+                    "lon": data["loc"].split(",")[1],
+                    "hostname": data["hostname"],
+                    "timezone": data["timezone"],
+                    "subnet": data["abuse"]["network"],
+                    "abuse": data["abuse"]["email"]
+                }
     except json.JSONDecodeError as e:
         logging.error("Failed to parse data for %s: %s", ip, e, exc_info=True)
         return HTTPException(500, "Failed to parse upstream response.")
@@ -126,21 +130,23 @@ async def _get_ipinfo_no_token(ip: str):
 
 async def _get_ipinfo_token(ip: str, token: str):
     try:
-        async with app.state.session.get("https://ipinfo.io/%s?token=%s" % (ip, token)) as response:
-            data = await response.json()
-            return {
-                "ip": data["ip"],
-                "city": data["city"],
-                "country": data["region"],
-                "countryCode": data["country"],
-                "asn": int(data["org"].split()[0][2:]),
-                "isp": data["org"].split(" ", 1)[1],
-                "source": "ipinfo",
-                "lat": data["loc"].split(",")[0],
-                "lon": data["loc"].split(",")[1],
-                "hostname": data["hostname"],
-                "timezone": data["timezone"],
-            }
+        logging.info("Querying IPInfo WITH token: %s", ip)
+        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(5)) as session:
+            async with session.get("https://ipinfo.io/%s?token=%s" % (ip, token)) as response:
+                data = await response.json()
+                return {
+                    "ip": data["ip"],
+                    "city": data["city"],
+                    "country": data["region"],
+                    "countryCode": data["country"],
+                    "asn": int(data["org"].split()[0][2:]),
+                    "isp": data["org"].split(" ", 1)[1],
+                    "source": "ipinfo",
+                    "lat": data["loc"].split(",")[0],
+                    "lon": data["loc"].split(",")[1],
+                    "hostname": data["hostname"],
+                    "timezone": data["timezone"],
+                }
     except json.JSONDecodeError as e:
         logging.error("Failed to parse data for %s: %s", ip, e, exc_info=True)
         return HTTPException(500, "Failed to parse upstream response.")
@@ -170,7 +176,7 @@ async def get_ip_information(ip: str) -> dict | HTTPException:
         logging.error("Failed to contact shronk net ip: %s", e, exc_info=True)
 
     try:
-        data = await get_ip_information(ip)
+        data = await get_from_ipinfo(ip)
     except Exception as e:
         logging.error("Failed to contact ipinfo: %s", e, exc_info=True)
     else:
@@ -195,7 +201,7 @@ async def ip(
         ip = ".".join(map(str, (random.randint(0, 255) for _ in range(4))))
 
     logging.info("looking up IP info for %s", ip)
-    data = await make_request(ip)
+    data = await get_ip_information(ip)
     if isinstance(data, HTTPException):
         raise data
     data["ip"] = ip