# Import stuff from datetime import datetime import json import urllib.request from urllib.parse import quote # Constants CT_HEADERS = { "X-CT-Version": "58380fc3", "X-CT-Timestamp": "1467966546", "X-CT-Client-Id": "cb7a42cd-efe2-42c3-a865-960396510735", } CT_CARDS = [ "Eurostar.FrequentTraveller", "Thalys.TheCard", "Thalys.ThePassBusiness", "Thalys.ThePassPremium", "Thalys.ThePassWeekEnd", "SNCF.Carte1225", "SNCF.CarteEscapades", "SNCF.CarteSenior", "SNCF.CarteEnfantPlus", "SNCF.CarteEnfantFamille", "SNCF.CarteFamilleNombreuse30", "SNCF.CarteFamilleNombreuse40", "SNCF.CarteFamilleNombreuse50", "SNCF.CarteFamilleNombreuse75", "SNCF.CarteMilitaireSecondClass", "SNCF.CarteMilitaireFirstClass", "SNCF.CarteFamilleMilitaire", "SNCF.AbonnementForfaitSecondClass", "SNCF.AbonnementForfaitFirstClass", "SNCF.AbonnementFrequenceSecondClass", "SNCF.AbonnementFrequenceFirstClass", "SNCF.CarteVoyageur", "SNCF.CarteGrandVoyageur", "SNCF.CarteGrandVoyageurPlus", "SNCF.CarteGrandVoyageurLeClub", "DB.BahnCard25SecondClass", "DB.BahnCard25FirstClass", "DB.BahnCard50SecondClass", "DB.BahnCard50FirstClass", "DB.BahnCard25BusinessSecondClass", "DB.BahnCard25BusinessFirstClass", "DB.BahnCard50BusinessSecondClass", "DB.BahnCard50BusinessFirstClass", "DB.BahnBonusCard", "Renfe.CarneJoven", "Renfe.TarjetaDorada", "Renfe.TarjetaMasRenfeJoven50", "Renfe.TarjetaMasRenfe", "Renfe.TarjetaMasRenfePlata", "Renfe.TarjetaMasRenfeOro", "Renfe.TarjetaMasRenfePremium", "Trenitalia.CartaFreccia", "Trenitalia.CartaVerde", "Trenitalia.CartaArgento", "NTV.ItaloPiu", "NS.VoordeelurenabonnementRailPlus", "NS.Voordeelurenabonnement", "CFF.DemiTarifRailPlus", "CFF.DemiTarif", "OBB.VORTEILScard", ] CT_SYSTEMS = [ "sncf", "db", "busbud", "idtgv", "ouigo", "trenitalia", "ntv", "hkx", "renfe", "timetable" ] # Structures class Station: def __init__(self, id, name, slug, info, latitude=0, longitude=0, score=1, is_sellable=True, parent_name=None, address=None): self._id = id self._name = name self._parent_name = parent_name @property def id(self): return self._id @property def name(self): return self._name @property def parent_name(self): return self._parent_name if self._parent_name is not None else self._name class ComfortClass: def __init__(self, id, name, cents, currency, condition_id, segment_id, default=False, is_available=False, options={}, reservation=None, title="", description=""): pass class Condition: def __init__(self, id, name, short_description="", comfort_class=None, segment_id=None): pass class Folder: def __init__(self, id, search_id, trip_ids, arrival_date, cents, travel_class, comfort, created_at, has_round_trip_fare, arrival_station_id, system, departure_date, is_sellable, is_birthdate_required, flexibility, digest, direction, currency, departure_station_id, is_phone_number_mandatory=False, is_only_possible_travel_class=False, has_fetchable_options=False, has_mandatory_fetchable_options=False, fetchable_options_hints=[], required_identification_documents=False, reference_folders=[]): pass class Passenger: def __init__(self, id, label): pass class Segment: def __init__(self, get_station, id, trip_id, train_number, digest, options, departure_date, train_name, arrival_station_id, comfort_class_ids, departure_station_id, travel_class, arrival_date, condition_id, carrier, train, co2_emission=0, reservation=False): self._id = id self._carrier = carrier self._train = train self._train_name = train_name self._train_number = train_number self._departure_date = parse_datetime(departure_date) self._departure_station = get_station(departure_station_id) self._arrival_date = parse_datetime(arrival_date) self._arrival_station = get_station(arrival_station_id) def __str__(self): return "%s %s %s" % (self.carrier, self.train_name, self.train_number) @property def id(self): return self._id @property def carrier(self): return self._carrier @property def train(self): return self._train @property def train_name(self): return self._train_name @property def train_number(self): return self._train_number @property def departure_date(self): return self._departure_date @property def departure_station(self): return self._departure_station @property def arrival_date(self): return self._arrival_date @property def arrival_station(self): return self._arrival_station class Trip: def __init__(self, get_station, get_segment, id, departure_date, arrival_date, cents, currency, folder_id, segment_ids, digest, departure_station_id, arrival_station_id, passenger_id): self._price = cents / 100 self._currency = currency self._digest = digest self._departure_date = parse_datetime(departure_date) self._arrival_date = parse_datetime(arrival_date) self._departure_station = get_station(departure_station_id) self._arrival_station = get_station(arrival_station_id) self._segments = [get_segment(segment_id) for segment_id in segment_ids] def __str__(self): via = [] segmnts = [] last_seg = None for seg in self.segments: if last_seg is not None: via.append(seg.departure_station.parent_name) if last_seg.arrival_station.id != seg.departure_station.id: segmnts.append("-> connection time: %s from %s to %s ->" % (seg.departure_date - last_seg.arrival_date, last_seg.arrival_station.name, seg.departure_station.name)) else: segmnts.append("-> connection time: %s in %s ->" % (seg.departure_date - last_seg.arrival_date, seg.departure_station.name)) segmnts.append(str(seg)) last_seg = seg return "From %s at %s to %s at %s%s // duration: %s // price: %s %s // %s" % ( self.departure_station.name, self.departure_date, self.arrival_station.name, self.arrival_date, " via " + ", ".join(via) if len(via) else "", self.duration, self.price, self.currency, " ".join(segmnts) ) @property def price(self): return self._price @property def currency(self): return self._currency @property def digest(self): return self._digest @property def departure_date(self): return self._departure_date @property def arrival_date(self): return self._arrival_date @property def departure_station(self): return self._departure_station @property def arrival_station(self): return self._arrival_station @property def duration(self): return self._arrival_date - self._departure_date @property def segments(self): return self._segments class Search: def __init__(self, comfort_classes, conditions, folders, passengers, search, segments, stations, trips): self._comfort_classes = [ComfortClass(**comfort_class) for comfort_class in comfort_classes] self._conditions = [Condition(**condition) for condition in conditions] self._folders = [Folder(**folder) for folder in folders] self._passengers = [Passenger(**passenger) for passenger in passengers] self._search = search self._stations = [Station(**station) for station in stations] self._segments = [Segment(self.get_station, **segment) for segment in segments] self._trips = [Trip(self.get_station, self.get_segment, **trip) for trip in trips] @property def trips(self): return self._trips @property def stations(self): return self._stations @property def segments(self): return self._segments def get_station(self, station_id): for station in self._stations: if station.id == station_id: return station def get_segment(self, segment_id): for segment in self._segments: if segment.id == segment_id: return segment # API calls def stations(name): """Get a list of stations matching given name""" headers = {"Accept": "application/json, text/javascript, */*; q=0.01"} headers.update(CT_HEADERS) req = urllib.request.Request('https://www.captaintrain.com/api/v5/stations?context=search&q=' + quote(name), headers=headers) with urllib.request.urlopen(req) as res: stations = json.loads(res.read().decode())["stations"] sorted(stations, key=lambda station: station["score"], reverse=True) for station in stations: yield Station(**station) def station(name): """Get the closest name matching station""" for station in stations(name): return station def search(departure, arrival, departure_date, return_date=None, via=None, age=27, systems=CT_SYSTEMS, cards=[]): "Search for a trip" headers = { "Accept": "application/json, text/javascript, */*; q=0.01", "Content-Type": "application/json; charset=utf-8", } headers.update(CT_HEADERS) req = urllib.request.Request('https://www.captaintrain.com/api/v5/search', headers=headers, data=json.dumps({ "search": { "departure_date": departure_date.strftime("%Y-%m-%dT%H:%M:%S%z") if departure_date else None, "return_date": return_date.strftime("%Y-%m-%dT%H:%M:%S%z") if return_date else None, "passengers": [ { "id": "d7c95386-9ed2-4366-8292-610d821940a3", "label": "adult", "age": age, "cards": [ { "reference": card } for card in cards ], "cui": None } ], "systems": systems if isinstance(systems, list) else [systems], "exchangeable_part": None, "departure_station_id": departure.id, "via_station_id": via.id if via is not None else None, "arrival_station_id": arrival.id, "exchangeable_pnr_id": None } }).encode()) with urllib.request.urlopen(req) as res: return Search(**json.loads(res.read().decode())) return None # Convinience functions def cheapest_trips(trips, max_price=None): """Return cheapest trips availables """ min_trips = [] for trip in trips: if max_price is not None and trip.price > max_price: continue elif len(min_trips) == 0 or min_trips[0].price > trip.price: min_trips = [trip] elif min_trips[0].price == trip.price and trip.digest not in map(lambda x: x.digest, min_trips): min_trips.append(trip) return min_trips def parse_datetime(s): try: return datetime.strptime(s, "%Y-%m-%dT%H:%M:%S%z") except ValueError: return datetime.strptime(s[::-1].replace(":", "", 1)[::-1], "%Y-%m-%dT%H:%M:%S%z")