From b251cd28e24c08e9170a11535fc121441b324d4f Mon Sep 17 00:00:00 2001
From: Ruoqi Yang <martinyangrq@outlook.com>
Date: Thu, 10 Nov 2022 12:29:10 +0800
Subject: add plugin: geoip loopup

---
 ATRI/configs/models.py |  6 ++++++
 ATRI/plugins/geoip.py  | 28 ++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+)
 create mode 100644 ATRI/plugins/geoip.py

diff --git a/ATRI/configs/models.py b/ATRI/configs/models.py
index a23d7f2..8e8869d 100644
--- a/ATRI/configs/models.py
+++ b/ATRI/configs/models.py
@@ -40,12 +40,18 @@ class Setu(BaseModel):
     reverse_proxy_domain: str
 
 
+class GeoIP(BaseModel):
+    account_id: str
+    license_key: str
+
+
 class ConfigModel(BaseModel):
     ConfigVersion: str
     BotConfig: BotConfig
     WithGoCQHTTP: WithGoCQHTTP
     SauceNAO: SauceNAO
     Setu: Setu
+    GeoIP: GeoIP
 
 
 class RuntimeConfig(BaseModel):
diff --git a/ATRI/plugins/geoip.py b/ATRI/plugins/geoip.py
new file mode 100644
index 0000000..a6521c2
--- /dev/null
+++ b/ATRI/plugins/geoip.py
@@ -0,0 +1,28 @@
+from nonebot.params import ArgStr
+import geoip2.webservice
+from ATRI.service import Service
+from ATRI import conf
+
+
+geoip = Service("GEOIP查询").document("search ip in MaxMind GEOIP databases")
+
+query_geoip = geoip.on_command("ip查询", "查询IP的地理位置", aliases={"IP查询", "查询IP"})
+
+LANG = "zh-CN"
+
+
+@query_geoip.got("ip_address", prompt="地址是?(支持ipv4/ipv6)")
+async def _(ip_address: str = ArgStr()):
+    with geoip2.webservice.Client(conf.GeoIP.account_id,
+                                  conf.GeoIP.license_key,
+                                  host="geolite.info") as client:
+        response = client.city(ip_address)
+        country = response.country.names[LANG]
+        city = response.city.names[LANG]
+        subdivision = ""
+        if response.subdivisions:
+            subdivision = response.subdivisions[0].names[LANG]
+        await query_geoip.finish(
+            f"IP: {ip_address}\n"
+            f"{country}{subdivision}{city}"
+        )
-- 
cgit v1.2.3