Add setup script, use Requests sessions, add README

README.md

@@ -0,0 +1,16 @@
+# VRR-Accountability
+A script to help evaluate the reliability of public transport.
+
+`monitor.py` logs departures (and their delays) to a MySQL database.
+
+## Requirements
+
+* A running MySQL database
+* Python 3 with the modules from `requirements.txt`
+    * These are installable with `pip3 install -r requirements.txt`
+
+## Usage
+
+* Run `./monitor.py -s` to generate a configuration file
+    * Store the output in `vrr.ini` and add your database configuration
+* Run `./monitor.py` to start logging!

monitor.py

@@ -8,6 +8,7 @@ import pause
 import sys
 import mysql.connector
 import configparser
+import argparse
 
 TABLE = """
 CREATE TABLE IF NOT EXISTS vrr (
@@ -32,41 +33,61 @@ class MOT:
     ALL_MODES = [LONG_DISTANCE_TRAIN, REGIONAL_TRAIN, COMMUTER_TRAIN, UNDERGROUND_TRAIN, TRAM, BUS, ELEVATED_TRAIN]
 
 
+parser = argparse.ArgumentParser()
+parser.add_argument("-s", "--setup", help="Run the setup routine.", action="store_true", dest="setup")
+
+args = parser.parse_args()
+if args.setup:
+    import setup
+    setup.setup()
+    sys.exit(0)
+
 # Parse the configuration file:
 
 cfg = configparser.ConfigParser()
-cfg.read('vrr.ini')
+try:
+    with open('vrr.ini') as f:
+        cfg.read_file(f)
+except IOError:
+    sys.exit("Could not open the configuration file.")
 
-db_config = {
+try:
+    db_config = {
         'user': cfg['db']['user'],
         'password': cfg['db']['pass'],
         'host': cfg['db']['host'],
         'database': cfg['db']['database'],
-}
+    }
 
-USE_MODES = []
+    USE_MODES = []
-if cfg['crawl']['use_long_distance']:
+    if cfg['crawl'].getboolean('use_long_distance'):
         USE_MODES.append(MOT.LONG_DISTANCE_TRAIN)
-if cfg['crawl']['use_regional_trains']:
+    if cfg['crawl'].getboolean('use_regional_trains'):
         USE_MODES.append(MOT.REGIONAL_TRAIN)
-if cfg['crawl']['use_commuter_trains']:
+    if cfg['crawl'].getboolean('use_commuter_trains'):
         USE_MODES.append(MOT.COMMUTER_TRAIN)
-if cfg['crawl']['use_trams']:
+    if cfg['crawl'].getboolean('use_trams'):
         USE_MODES.append(MOT.TRAM)
-if cfg['crawl']['use_buses']:
+    if cfg['crawl'].getboolean('use_buses'):
         USE_MODES.append(MOT.BUS)
-if cfg['crawl']['use_elevated_trains']:
+    if cfg['crawl'].getboolean('use_elevated_trains'):
         USE_MODES.append(MOT.ELEVATED_TRAIN)
 
-if cfg['crawl']['station_id'] is not None:
+    if cfg['crawl']['station_id'] is not None:
-    USE_STATION_ID = cfg['crawl']['station_id']
+        USE_STATION_ID = cfg['crawl'].getint('station_id')
-else:
+    else:
         sys.exit("Please specify a station_id in the [crawl] section of vrr.ini")
-USE_LINES = cfg['crawl']['use_lines'].split(',')
+    USE_LINES = cfg['crawl']['use_lines'].split(',')
+except (IndexError, configparser.NoOptionError, configparser.NoSectionError):
+    sys.exit("There is something wrong with the configuration file. Exiting.")
 
 ALL_LINES = []
 TRIP_CANCELLED = -9999
 
+# Initialize Requests session
+
+HTTP = requests.session()
+
 
 def make_request_data(station_id: int, result_count: int = 8, modes: List = MOT.ALL_MODES,
                       lines: List[str] = ALL_LINES) -> dict:
@@ -122,7 +143,7 @@ def make_request_data(station_id: int, result_count: int = 8, modes: List = MOT.
 
 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 = HTTP.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()

setup.py

@@ -0,0 +1,115 @@
+import requests
+import sys
+from typing import List
+import configparser
+
+HTTP = requests.session()
+
+
+def yn(s: str) -> bool:
+    return s in ['y', 'Y', 'J', 'j', '']
+
+
+def search_station(search: str) -> List or None:
+    params = {
+        'query': search.replace(" ", "+")
+    }
+    resp = HTTP.get("https://abfahrtsmonitor.vrr.de/backend/api/stations/search", params=params)
+    try:
+        resp.raise_for_status()
+    except requests.exceptions.HTTPError:
+        return None
+
+    return resp.json()['suggestions']
+
+
+def get_station() -> int:
+    station_id = None
+    while station_id is None:
+        search = input("Which station would you like to monitor? ")
+        print("Getting suggestions...")
+        results = search_station(search)
+        if results:  # empty lists and None are False
+            for ptr in range(len(results)):
+                print(str(ptr) + ". " + results[ptr]['value'] + "\t" + results[ptr]['data'])
+            choice_ptr = None
+            if len(results) > 1:
+                while choice_ptr is None:
+                    parse = input("Which one? ")
+                    try:
+                        parse = int(parse)
+                        test = results[parse]
+                    except ValueError:
+                        print("You did not input a number!")
+                    except IndexError:
+                        print("Number out of range.")
+                    choice_ptr = parse
+            else:
+                choice_ptr = 0
+            station_id = results[choice_ptr]['data']
+        else:
+            print('Got no results for "' + search + '".')
+    return int(station_id)
+
+
+def get_lines(station_id: int) -> List[str]:
+    print("Getting lines...")
+    resp = HTTP.get("https://abfahrtsmonitor.vrr.de/backend/api/lines/" + str(station_id) + "/search")
+    resp.raise_for_status()
+    results = resp.json()
+    choices = ""
+    if len(results) > 20:
+        page = yn(input("There are more than 20 results. Would you like to view them page-by-page? (Y/n)"))
+    else:
+        page = False
+    if page:
+        print("Paging in batches of 20 results.")
+        results_displayed = 0
+        for ptr in range(len(results)):
+            if ptr % 20 == 0 or ptr == len(results) - 1 and ptr != 0:
+                choices += input("Please input your choices as a space-separated list (e.g. '0 2 7 15'):\n")
+            print(str(ptr) + ". " + results[ptr]['name'])
+    else:
+        for ptr in range(len(results)):
+            print(str(ptr) + ". " + results[ptr]['name'])
+        choices += input("Please input your choices as a space-separated list (e.g. '0 2 7 15'):\n") + ' '
+    filt_arr = []
+    for ptr in choices.split(" "):
+        if ptr == '':
+            continue
+        try:
+            filt_arr.append(results[int(ptr)])
+        except ValueError:
+            print(ptr + " is not a number")
+        except IndexError:
+            print(ptr + " is out of range")
+    return ["{0}:{1}:{2}:{3}".format(r['network'], r['line'], r['supplement'], r['directionCode'])
+            for r in filt_arr]
+
+
+def setup() -> None:
+    station_id = get_station()
+    lines_ch = input("Would you like to choose specific lines? (Y/n)", )
+    if yn(lines_ch):
+        lines = get_lines(station_id)
+    else:
+        lines = None
+    cfg = configparser.ConfigParser()
+    cfg.add_section('crawl')
+    cfg.add_section('db')
+    cfg['crawl']['station_id'] = str(station_id)
+    cfg['crawl']['use_long_distance'] = 'yes'
+    cfg['crawl']['use_regional_trains'] = 'yes'
+    cfg['crawl']['use_commuter_trains'] = 'yes'
+    cfg['crawl']['use_underground_trains'] = 'yes'
+    cfg['crawl']['use_trams'] = 'yes'
+    cfg['crawl']['use_buses'] = 'yes'
+    cfg['crawl']['use_elevated_trains'] = 'yes'
+    cfg['crawl']['use_lines'] = ",".join(lines) if lines is not None else ""
+    cfg['db']['user'] = 'vrr'
+    cfg['db']['pass'] = 'vrr'
+    cfg['db']['host'] = 'localhost'
+    cfg['db']['database'] = 'vrr'
+    print("Please save the following output to 'vrr.ini' and adjust any further settings:")
+    print("\n" * 3)
+    cfg.write(sys.stdout)

vrr.ini

@@ -1,3 +1,8 @@
+; =====================================
+; VRR_ACCOUNTABILITY CONFIGURATION FILE
+; =====================================
+; You may generate a configuration file using ./monitor.py -s
+
 [crawl]
 ; A valid EFA station ID. This can be obtained in the Öffi "departures" view or by inspecting
 ; the regular POST requests on abfahrtsmonitor.vrr.de.
@@ -12,7 +17,7 @@ use_buses=yes
 use_elevated_trains=yes
 ; The server allows filtering by line and direction. These line names can be
 ; obtained by inspecting the POST requests on abfahrtsmonitor.vrr.de.
-use_lines=rbg:70070:H,rbg:70070:R,rbg:70076:H,rbg:70076:H
+use_lines=rbg:70070: :H,rbg:70070: :R,rbg:70076: :H,rbg:70076: :H
 
 [db]
 user=vrr