@nomadamas/k-skill 0.2.5 → 0.4.0

package/skills/srt-booking/scripts/srt_booking.py

@@ -1,272 +1,171 @@
-#!/usr/bin/env python3
+#!/usr/bin/env -S uv run --locked --script
+# /// script
+# requires-python = ">=3.11"
+# dependencies = ["SRTrain==2.6.7"]
+# ///
+"""Live, anonymous, read-only SRT timetable lookup through SRTrain."""
+
 from __future__ import annotations
 
 import argparse
-import base64
-import contextlib
-import io
-import importlib
 import json
-import os
+import re
 import sys
-from types import ModuleType
-from typing import Protocol
-
-from srt_seats import parse_cars, parse_seats, sort_cars_for_booking, sort_seats_for_booking
-
-SRT_SEAT_ENDPOINT = "https://etk.srail.kr/hpg/hra/01/selectPassengerResearchList.do"
-TRAIN_ID_PREFIX = "srt:v1:"
-TRAIN_ID_FIELDS = (
-    "train_number",
-    "dep_date",
-    "dep_time",
-    "arr_date",
-    "arr_time",
-    "train_code",
-    "dep_station_code",
-    "arr_station_code",
-    "dep_station_run_order",
-    "arr_station_run_order",
-)
-ROOM_CODE = {"general": "1", "special": "2"}
-ROOM_NAME = {"general": "일반실", "special": "특실"}
-
-
-class SrtTrainLike(Protocol):
-    train_number: str
-    dep_date: str
-    dep_time: str
-    arr_date: str
-    arr_time: str
-    train_code: str
-    train_name: str
-    dep_station_code: str
-    dep_station_name: str
-    arr_station_code: str
-    arr_station_name: str
-    dep_station_run_order: str
-    arr_station_run_order: str
-    general_seat_state: str
-    special_seat_state: str
-    reserve_wait_possible_code: str
-
-    def general_seat_available(self) -> bool: ...
-
-    def special_seat_available(self) -> bool: ...
-
-    def reserve_standby_available(self) -> bool: ...
-
-
-class ResponseLike(Protocol):
-    text: str
-
-
-class SessionLike(Protocol):
-    def get(self, url: str, params: dict[str, str]) -> ResponseLike: ...
-
-
-class SrtClientLike(Protocol):
-    _session: SessionLike
-
-    def search_train(
-        self,
-        dep: str,
-        arr: str,
-        date: str,
-        time: str,
-        time_limit: str | None = None,
-        available_only: bool = True,
-    ) -> list[SrtTrainLike]: ...
-
-
-def load_srt_module() -> ModuleType:
-    try:
-        return importlib.import_module("SRT")
-    except ModuleNotFoundError as exc:
-        raise SystemExit("scripts/srt_booking.py requires SRTrain: python3 -m pip install SRTrain")
+from contextlib import redirect_stdout
+from datetime import date as calendar_date
+from datetime import time
+from typing import Any
 
+from requests import RequestException
+from SRT import SRT
+from SRT.constants import API_ENDPOINTS, STATION_CODE
+from SRT.errors import SRTError, SRTNetFunnelError
+from SRT.netfunnel import NetFunnelHelper
 
-def build_client(auto_login: bool = False) -> SrtClientLike:
-    srt_module = load_srt_module()
-    srt_id = os.environ.get("KSKILL_SRT_ID", "")
-    srt_pw = os.environ.get("KSKILL_SRT_PASSWORD", "")
-    return srt_module.SRT(srt_id, srt_pw, auto_login=auto_login)
+BOOKING_URL = "https://etk.srail.kr/hpg/hra/01/selectScheduleList.do?pageId=TK0101010000"
+NETFUNNEL_URL = "https://nf.letskorail.com/ts.wseq"
 
 
-def train_id_payload(train: SrtTrainLike) -> dict[str, str]:
-    return {field: getattr(train, field) for field in TRAIN_ID_FIELDS}
+def build_client() -> SRT:
+    NetFunnelHelper.NETFUNNEL_URL = NETFUNNEL_URL
+    return SRT("", "", auto_login=False)
 
 
-def build_train_id(train: SrtTrainLike) -> str:
-    raw = json.dumps(train_id_payload(train), ensure_ascii=False, separators=(",", ":")).encode()
-    encoded = base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")
-    return f"{TRAIN_ID_PREFIX}{encoded}"
+def source_info() -> dict[str, str]:
+    return {
+        "mode": "live",
+        "transport": "SRTrain",
+        "operator": "주식회사 에스알",
+        "endpoint": API_ENDPOINTS["search_schedule"],
+        "queue_endpoint": NETFUNNEL_URL,
+        "authentication": "anonymous",
+        "mutation": "none; timetable search only",
+        "booking_url": BOOKING_URL,
+    }
 
 
-def parse_train_id(train_id: str) -> dict[str, str]:
-    if not train_id.startswith(TRAIN_ID_PREFIX):
-        raise SystemExit("train_id must start with srt:v1:")
-    encoded = train_id.removeprefix(TRAIN_ID_PREFIX)
-    padded = encoded + ("=" * ((4 - len(encoded) % 4) % 4))
-    try:
-        payload = json.loads(base64.urlsafe_b64decode(padded.encode()).decode())
-    except (ValueError, json.JSONDecodeError, UnicodeDecodeError) as exc:
-        raise SystemExit("train_id is invalid; rerun search and copy a fresh train_id") from exc
-    if not isinstance(payload, dict):
-        raise SystemExit("train_id is invalid; rerun search and copy a fresh train_id")
-    if any(not isinstance(payload.get(field), str) or not payload[field] for field in TRAIN_ID_FIELDS):
-        raise SystemExit("train_id is invalid; rerun search and copy a fresh train_id")
-    return {field: payload[field] for field in TRAIN_ID_FIELDS}
+def format_time(value: str) -> str:
+    return f"{value[:2]}:{value[2:4]}"
 
 
-def find_train_by_id(trains: list[SrtTrainLike], train_id: str) -> SrtTrainLike | None:
-    expected = parse_train_id(train_id)
-    return next((train for train in trains if train_id_payload(train) == expected), None)
+def normalize_station(value: str) -> str:
+    """Accept common "...역" input for a canonical SRT station name."""
+    name = value.strip()
+    if name in STATION_CODE:
+        return name
+    if name.endswith("역") and name[:-1] in STATION_CODE:
+        return name[:-1]
+    return name
 
 
-def normalize_train(train: SrtTrainLike, index: int) -> dict[str, str | bool | int]:
+def normalize_train(train: Any) -> dict[str, Any]:
     return {
-        "index": index,
-        "train_id": build_train_id(train),
-        "train_no": train.train_number,
-        "train_type": train.train_name,
-        "dep_name": train.dep_station_name,
-        "dep_date": train.dep_date,
-        "dep_time": train.dep_time,
-        "arr_name": train.arr_station_name,
-        "arr_date": train.arr_date,
-        "arr_time": train.arr_time,
-        "has_general_seat": train.general_seat_available(),
-        "has_special_seat": train.special_seat_available(),
-        "has_waiting_list": train.reserve_standby_available(),
+        "train_no": str(train.train_number),
+        "train_type": str(train.train_name),
+        "dep": str(train.dep_station_name),
+        "arr": str(train.arr_station_name),
+        "dep_date": str(train.dep_date),
+        "dep_time": format_time(str(train.dep_time)),
+        "arr_time": format_time(str(train.arr_time)),
+        "general_seat_available": bool(train.general_seat_available()),
+        "special_seat_available": bool(train.special_seat_available()),
     }
 
 
-def seat_page_params(train: SrtTrainLike, room: str, car_no: int | None) -> dict[str, str]:
+def search_live_timetable(
+    *,
+    dep: str,
+    arr: str,
+    date: str,
+    earliest: str,
+    latest: str,
+    limit: int,
+) -> dict[str, Any]:
+    validate_date(date)
+    start = validate_time(earliest)
+    end = validate_time(latest)
+    if start > end:
+        raise ValueError("--time must not be later than --time-limit")
+    client = build_client()
+    with redirect_stdout(sys.stderr):
+        trains = client.search_train(
+            dep=normalize_station(dep),
+            arr=normalize_station(arr),
+            date=date,
+            time=start,
+            time_limit=end,
+            available_only=False,
+        )
     return {
-        "runDt1": train.dep_date,
-        "dptDt1": train.dep_date,
-        "dptTm1": train.dep_time,
-        "trnNo1": f"{int(train.train_number):05d}",
-        "trnGpCd1": "300",
-        "dptRsStnCd1": train.dep_station_code,
-        "arvRsStnCd1": train.arr_station_code,
-        "dptStnRunOrdr1": train.dep_station_run_order,
-        "arvStnRunOrdr1": train.arr_station_run_order,
-        "seatAttCd1": "015",
-        "psrmClCd1": ROOM_CODE[room],
-        "index1": "0",
-        "scarNo1": "" if car_no is None else f"{car_no:04d}",
-        "chtnDvCd": "1",
-        "jrnySqno": "001",
-        "mode": "1",
-        "psgNum": "1",
-        "pageId": "",
+        "count": min(len(trains), limit),
+        "trains": [normalize_train(train) for train in trains[:limit]],
+        "date": date,
+        "source": source_info(),
+        "schedule_note": "실시간 시간표·좌석 가능 여부 조회이며 예약·좌석 선점은 실행하지 않습니다.",
+        "booking_url": BOOKING_URL,
     }
 
 
-def fetch_seat_page(client: SrtClientLike, train: SrtTrainLike, room: str, car_no: int | None) -> str:
-    with contextlib.redirect_stdout(io.StringIO()):
-        response = client._session.get(SRT_SEAT_ENDPOINT, params=seat_page_params(train, room, car_no))
-    return response.text
-
-
-def command_search(args: argparse.Namespace) -> None:
-    client = build_client(auto_login=False)
-    with contextlib.redirect_stdout(io.StringIO()):
-        trains = client.search_train(args.dep, args.arr, args.date, args.time, args.time_limit, args.available_only)
-    print_json({"count": len(trains[: args.limit]), "trains": [normalize_train(train, index) for index, train in enumerate(trains[: args.limit], 1)]})
-
-
-def command_seats(args: argparse.Namespace) -> None:
-    client = build_client(auto_login=False)
-    with contextlib.redirect_stdout(io.StringIO()):
-        trains = client.search_train(args.dep, args.arr, args.date, args.time, args.time_limit, available_only=False)
-    train = find_train_by_id(trains, args.train_id)
-    if train is None:
-        raise SystemExit("train_id no longer matches any current search result; rerun search and choose a fresh train_id")
-
-    initial_html = fetch_seat_page(client, train, args.room, args.car_no)
-    cars = [car for car in parse_cars(initial_html) if car["room_class"] == ROOM_NAME[args.room]]
-    if args.car_no is not None:
-        cars = [car for car in cars if car["car_no"] == args.car_no]
-    else:
-        cars = [car for car in cars if car["available"]]
-    if not cars:
-        raise SystemExit(f"seat car data is unavailable for {args.room}; retry search or choose another train")
-
-    car_payloads: list[dict[str, object]] = []
-    for car in sort_cars_for_booking(cars, args.car_priority):
-        html = initial_html if args.car_no == car["car_no"] else fetch_seat_page(client, train, args.room, car["car_no"])
-        seats = parse_seats(html)
-        if args.seat:
-            seats = [seat for seat in seats if seat["seat"] == args.seat]
-        seats = sort_seats_for_booking(seats, args.seat_priority)
-        if args.available_only:
-            seats = [seat for seat in seats if seat["available"]]
-        available_seats = [seat for seat in seats if seat["available"]]
-        limited = seats[: args.limit]
-        payload = dict(car)
-        payload["available_seat_count"] = len(available_seats)
-        payload["available_seats"] = [seat["seat"] for seat in available_seats]
-        payload["shown_seat_count"] = len(limited)
-        payload["seats"] = limited
-        if args.seat:
-            payload["requested_seat"] = args.seat
-            payload["requested_seat_available"] = any(seat["available"] for seat in seats)
-        car_payloads.append(payload)
-
-    print_json({
-        "train": normalize_train(train, 1),
-        "room": args.room,
-        "available_only": args.available_only,
-        "car_priority": args.car_priority,
-        "seat_priority": args.seat_priority,
-        "cars": car_payloads,
-    })
+def validate_date(value: str) -> str:
+    if not re.fullmatch(r"\d{8}", value):
+        raise ValueError("date must use YYYYMMDD")
+    try:
+        calendar_date(int(value[:4]), int(value[4:6]), int(value[6:]))
+    except ValueError as exc:
+        raise ValueError("date must use a valid YYYYMMDD value") from exc
+    return value
 
 
-def print_json(payload: dict[str, object]) -> None:
-    print(json.dumps(payload, ensure_ascii=False, indent=2))
+def validate_time(value: str) -> str:
+    if not re.fullmatch(r"\d{4}", value):
+        raise ValueError("time must use HHMM")
+    try:
+        time.fromisoformat(f"{value[:2]}:{value[2:]}")
+    except ValueError as exc:
+        raise ValueError("time must use a valid HHMM value") from exc
+    return value + "00"
 
 
 def build_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(description="SRT search and seat lookup helper for k-skill")
-    subparsers = parser.add_subparsers(dest="command", required=True)
-    search = subparsers.add_parser("search", help="SRT 열차를 조회합니다")
-    add_trip_args(search)
-    search.add_argument("--time-limit", default=None, help="조회 종료 시각 HHMMSS")
-    search.add_argument("--available-only", action="store_true", default=False, help="예약 가능한 열차만 출력")
-    search.add_argument("--limit", type=int, default=5, help="출력할 최대 열차 수")
-    search.set_defaults(func=command_search)
-
-    seats = subparsers.add_parser("seats", help="SRT 호차별 좌석번호를 조회합니다")
-    add_trip_args(seats)
-    seats.add_argument("--train-id", required=True, help="search 결과에서 복사한 stable train_id")
-    seats.add_argument("--time-limit", default=None, help="조회 종료 시각 HHMMSS")
-    seats.add_argument("--room", choices=sorted(ROOM_CODE), default="general")
-    seats.add_argument("--car-no", type=int, default=None, help="특정 호차만 조회")
-    seats.add_argument("--seat", default=None, help="특정 좌석번호만 조회, 예: 6C")
-    seats.add_argument("--available-only", action="store_true", help="빈 좌석만 출력")
-    seats.add_argument("--car-priority", choices=("center", "low", "high"), default="center")
-    seats.add_argument("--seat-priority", choices=("forward-window", "window-forward", "row-low"), default="forward-window")
-    seats.add_argument("--limit", type=int, default=100, help="호차별 출력할 최대 좌석 수")
-    seats.set_defaults(func=command_seats)
+    parser = argparse.ArgumentParser(description="SRT live timetable lookup through SRTrain (read-only)")
+    commands = parser.add_subparsers(dest="command", required=True)
+    search = commands.add_parser("search", help="query the current SRT timetable")
+    search.add_argument("--dep", required=True)
+    search.add_argument("--arr", required=True)
+    search.add_argument("--date", required=True, help="YYYYMMDD")
+    search.add_argument("--time", default="0000", help="earliest departure, HHMM")
+    search.add_argument("--time-limit", default="2359", help="latest departure, HHMM")
+    search.add_argument("--limit", type=int, default=10)
+    commands.add_parser("source", help="show the read-only live query endpoint")
     return parser
 
 
-def add_trip_args(parser: argparse.ArgumentParser) -> None:
-    parser.add_argument("dep", help="출발역")
-    parser.add_argument("arr", help="도착역")
-    parser.add_argument("date", help="출발일 YYYYMMDD")
-    parser.add_argument("time", help="희망 시작 시각 HHMMSS")
-
-
 def main(argv: list[str] | None = None) -> int:
-    args = build_parser().parse_args(argv)
-    args.func(args)
-    return 0
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    try:
+        if args.command == "source":
+            print(json.dumps(source_info(), ensure_ascii=False, indent=2))
+            return 0
+        if args.limit < 1 or args.limit > 50:
+            raise ValueError("--limit must be between 1 and 50")
+        result = search_live_timetable(
+            dep=args.dep,
+            arr=args.arr,
+            date=args.date,
+            earliest=args.time,
+            latest=args.time_limit,
+            limit=args.limit,
+        )
+        print(json.dumps(result, ensure_ascii=False, indent=2))
+        return 0
+    except (SRTNetFunnelError, RequestException):
+        parser.error(f"SRT timetable lookup unavailable; use the official page: {BOOKING_URL}")
+        return 2
+    except (SRTError, ValueError) as exc:
+        parser.error(str(exc))
+        return 2
 
 
 if __name__ == "__main__":
-    raise SystemExit(main(sys.argv[1:]))
+    sys.exit(main())