From e3c2a76cc7907a7da9565caa19ccb957c3b8300a Mon Sep 17 00:00:00 2001
From: XieXiLin <support@xiexilin.com>
Date: Sat, 20 Jan 2024 02:33:45 +0800
Subject: [PATCH] add maintancer & alias

---
 nonebot_plugin_cnrail/data_source.py          | 36 +++++++++++++++++--
 nonebot_plugin_cnrail/models.py               |  1 +
 .../templates/template.html.jinja             |  6 ++--
 3 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/nonebot_plugin_cnrail/data_source.py b/nonebot_plugin_cnrail/data_source.py
index ebc92e1..267ff86 100644
--- a/nonebot_plugin_cnrail/data_source.py
+++ b/nonebot_plugin_cnrail/data_source.py
@@ -13,6 +13,8 @@
 CHINA_RAIL_SEARCH_API = "https://search.12306.cn/search/v1/train/search"
 CHINA_RAIL_DETAIL_API = "https://kyfw.12306.cn/otn/queryTrainInfo/query"
 
+CNRAIL_DATA_BASE_URL = "https://cnrail-data.baka.pub/data/"
+
 ACG_IMAGE_URL = "https://www.loliapi.com/acg/pe/"
 
 TEMPLATE_PATH = Path(__file__).parent / "templates" / "template.html.jinja"
@@ -42,7 +44,28 @@ async def query_train_info(train_code: str, train_date: str) -> Optional[TrainIn
 
     raw_data = resp.json()["data"]
     if not raw_data:
-        return None
+        async with httpx.AsyncClient(
+            base_url=CNRAIL_DATA_BASE_URL,
+            follow_redirects=True,
+        ) as client:
+            resp = await client.get("/alias.json")
+            resp.raise_for_status()
+
+        train_code = resp.json()[train_code]
+
+        async with httpx.AsyncClient(follow_redirects=True) as client:
+            resp = await client.get(
+                CHINA_RAIL_SEARCH_API,
+                params={
+                    "keyword": train_code,
+                    "date": train_date.replace("-", ""),
+                },
+            )
+            resp.raise_for_status()
+
+        raw_data = resp.json()["data"]
+        if not raw_data:
+            return None
 
     data = parse_obj_as(List[TrainSummary], raw_data)
     if len(data) > 1:
@@ -74,7 +97,16 @@ async def query_train_info(train_code: str, train_date: str) -> Optional[TrainIn
         return None
 
     stations = parse_obj_as(List[TrainStation], raw_data)
-    return TrainInfo(summary=summary, stations=stations)
+
+    async with httpx.AsyncClient(
+        base_url=CNRAIL_DATA_BASE_URL, follow_redirects=True
+    ) as client:
+        resp = await client.get("/maintance.json")
+        resp.raise_for_status()
+
+    maintancer = resp.json()[train_code]
+
+    return TrainInfo(summary=summary, stations=stations, maintancer=maintancer)
 
 
 async def render_train_info(info: TrainInfo) -> bytes:
diff --git a/nonebot_plugin_cnrail/models.py b/nonebot_plugin_cnrail/models.py
index 856aa15..29d59a8 100644
--- a/nonebot_plugin_cnrail/models.py
+++ b/nonebot_plugin_cnrail/models.py
@@ -49,6 +49,7 @@ def stay_minutes(self) -> int:
 class TrainInfo(BaseModel):
     summary: TrainSummary
     stations: List[TrainStation]
+    maintancer: str
 
     _arrive_next_day: Optional[bool] = PrivateAttr(None)
 
diff --git a/nonebot_plugin_cnrail/templates/template.html.jinja b/nonebot_plugin_cnrail/templates/template.html.jinja
index 65075dd..76467b7 100644
--- a/nonebot_plugin_cnrail/templates/template.html.jinja
+++ b/nonebot_plugin_cnrail/templates/template.html.jinja
@@ -320,9 +320,9 @@
               </div>
             </div>
           </div>
-          <!-- <div class="station-bottom">
-            <div class="admin">担当路局：郑州局</div>
-          </div> -->
+          <div class="station-bottom">
+            <div class="admin">担当路局：{{ info.maintancer }}</div>
+          </div>
         </div>
       </div>