vrr-accountability/monitor.py

#!/usr/bin/env python3
import requests
import json
from typing import List
from pprint import pprint, pformat
import datetime
import pause
import sys


class MOT:
    LONG_DISTANCE_TRAIN = 0
    REGIONAL_TRAIN = 1
    COMMUTER_TRAIN = 2
    UNDERGROUND_TRAIN = 3
    TRAM = 4
    BUS = 15
    ELEVATED_TRAIN = 6
    ALL_MODES = [LONG_DISTANCE_TRAIN, REGIONAL_TRAIN, COMMUTER_TRAIN, UNDERGROUND_TRAIN, TRAM, BUS, ELEVATED_TRAIN]


ALL_LINES = []
TRIP_CANCELLED = -9999

lines_filter = [
    'rbg:70070: :H',  # U70 -> Düsseldorf Hbf
    'rbg:70070: :R',  # U70 -> Krefeld Rheinstr
    'rbg:70076: :H',  # U76 -> Düsseldorf Hbf
    'rbg:70076: :R',  # U76 -> Krefeld Rheinstr
]


def t(s: str) -> str:
    """
    Encode a string to be used as a station identifier.
    :param s: a string to encode
    :return: the encoded string
    """
    return s.replace(' ', '+')


def make_request_data(station_id: int, result_count: int = 8, modes: List = MOT.ALL_MODES, lines: List[str] = ALL_LINES) -> dict:
    """
    Prepare a request data dictionary to put into get_data()
    :param station_id: an EFA station ID
    :param result_count: how many departures to return
    :param modes: which modes of transport to use
    :param lines: which lines to use (line identifiers look like 'provider:line ID: :direction ID',
    e.g. 'rbg:70070: :H' for the Rheinbahn U70 to Düsseldorf Hbf.
    :return: a dictionary with the data necessary to make a request to the Abfahrtsmonitor API.
    """
    """
    The request data dictionary can have the following items:
    
    stationID: a numerical EFA station ID
    stationName: (optional) the station's name
    platformVisibility: (optional) ???
    transport: a comma-separated list of the modes of transport to be displayed. See the constants for values.
    useAllLines: display all available lines or filter them using the linesFilter
    linesFilter: a JSON array with the lines to be displayed. See lines_filter for the format
    optimizedForStation: (optional) ???
    rowCount: the amount of results to be returned
    refreshInterval: (optional) (display parameter) refresh rate in seconds for the browser UI
    distance: (optional) (display parameter) distance from the monitor to the stop
    marquee: (optional) (display parameter) make the path text scroll sideways
    sortBy: (optional) ???
    """
    request_data = {
        'stationId': int(station_id),
        'rowCount': result_count
    }

    # sanity check: do the modes exist?
    for mode in modes:
        if mode not in MOT.ALL_MODES:
            raise ValueError(str(mode) + "Unknown transport mode!")

    # Add the list to the data dictionary
    request_data['transport'] = ','.join("{0}".format(n) for n in modes).rstrip(',')

    if lines is ALL_LINES:
        request_data['useAllLines'] = 1
    else:
        lines_dictarr = [{'data': t(v)} for v in lines]
        request_data['linesFilter'] = json.dumps(lines_dictarr)
        request_data['useAllLines'] = 0

    # finally, add the HTML naming
    request_data = {"table[departure][{0}]".format(k): v for k, v in request_data.items()}
    return request_data


def get_data(request_data: dict, headers: dict = None, cookies: dict = None) -> dict:
    url = 'https://abfahrtsmonitor.vrr.de/backend/api/stations/table'
    reply = requests.post(url, data=request_data, headers=headers, cookies=cookies)
    reply.raise_for_status()
    print('Request time elapsed: ' + str(reply.elapsed), file=sys.stderr)
    return reply.json()


def is_cancelled(trip: dict) -> bool:
    if trip['delay'] is not None:
        return int(trip['delay']) == TRIP_CANCELLED
    return False


def is_late(trip: dict) -> bool:
    if trip['delay'] is not None:
        return int(trip['delay']) > 0
    return False


def is_early(trip: dict) -> bool:
    if trip['delay'] is not None:
        return int(trip['delay']) < 0 and int(trip['delay']) != TRIP_CANCELLED
    return False


# Pretty-print the reply data.
"""print("Data:")
pprint(reply_data)"""


def fixup_data(d: dict) -> dict:
    for trip in d['departureData']:
        if trip['delay'] == '':
            trip['delay'] = None
    return d


def print_trip(trip: dict) -> None:
    trip_part = "The {}:{} {} (???:{}: :{}) service to {} ".format(trip['hour'], trip['minute'], trip['lineNumber'], trip['lineCode'], trip['directionCode'], trip['direction'])
    if is_cancelled(trip):
        print(trip_part + "is cancelled.")
    elif is_late(trip):
        print(trip_part + "is {} minutes late.".format(trip['delay']))
    elif is_early(trip):
        print(trip_part + "is {} minutes early.".format(-trip['delay']))


def get_next_refresh(data: dict):
    times = []
    for trip in data['departureData']:
        times.append(trip['orgFullTime'])
        times.append(trip['fullTime'])
    times = [int(time) for time in times if int(time) > datetime.datetime.now().timestamp()]
    times.sort()
    for time in times:
        if (datetime.datetime.fromtimestamp(time) - datetime.datetime.now()) > datetime.timedelta(seconds=30):
            return time
    return (datetime.datetime.now() + datetime.timedelta(seconds=60)).timestamp()


def update():
    reply_data = get_data(
        make_request_data(
            20021002,
            8,
            lines=lines_filter
        )
    )
    reply_data = fixup_data(reply_data)
    for trip in reply_data['departureData']:
        print_trip(trip)
    return reply_data


def wait():
    data = update()
    while True:
        next_refresh = get_next_refresh(data)
        print("Sleeping until " + datetime.datetime.fromtimestamp(next_refresh).isoformat(), file=sys.stderr)
        pause.until(next_refresh)
        data = update()


def main():
    wait()


main()