#!/usr/bin/env python3 from __future__ import annotations import argparse import base64 import contextlib import io import importlib import json import os 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") 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) def train_id_payload(train: SrtTrainLike) -> dict[str, str]: return {field: getattr(train, field) for field in TRAIN_ID_FIELDS} 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 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 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_train(train: SrtTrainLike, index: int) -> dict[str, str | bool | int]: 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(), } def seat_page_params(train: SrtTrainLike, room: str, car_no: int | None) -> dict[str, str]: 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": "", } 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 print_json(payload: dict[str, object]) -> None: print(json.dumps(payload, ensure_ascii=False, indent=2)) 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) 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 if __name__ == "__main__": raise SystemExit(main(sys.argv[1:]))