New SNCF module

main.py

@@ -81,8 +81,10 @@ def main():
     # alerts
     alerts = []
 
-    from modules.weather import WeatherAlerts
     alerts += [a for a in RATPWeatherModule().gen_alerts()]
+    from modules.sncf import SNCFWeatherModule
+    alerts += SNCFWeatherModule().gen_alerts("normandie")
+    from modules.weather import WeatherAlerts
     alerts += WeatherAlerts().gen_alerts()
 
     from modules import AlertsModule

modules/sncf.py

@@ -0,0 +1,101 @@
+from datetime import datetime, timedelta, timezone
+import base64
+import json
+import os
+import urllib.parse
+import urllib.request
+import re
+
+from PIL import Image, ImageDraw, ImageFont, ImageOps
+
+class SNCFAPI:
+
+    CLEANR = re.compile('<.*?>')
+
+    def __init__(self):
+        self.baseurl = "https://www.sncf.com/api/iv"
+        self.auth = base64.b64encode(b"admin:$2y$10$QvxWSS4f5DIFSkAmuBoV9OJG3M2bhec9d3F2.YnULBMtpzKAq2KS.").decode()
+
+        self._cached_file = ".sncf-%d-%s.cache"
+        self.cache_time = 5
+
+    def get_icon(size):
+        height = int(size * 0.531)
+        img = Image.open("icons/sncf.png").resize((size, height))
+        r,g,b,a = img.split()
+        rgb_image = Image.merge('RGB', (r,g,b))
+        inverted_image = ImageOps.invert(rgb_image)
+        r2,g2,b2 = inverted_image.split()
+
+        return Image.merge('RGBA', (r2,g2,b2,a))
+
+    def get_train_status(self, numero, date):
+        cache_file = self._cached_file % (numero, date.strftime("%Y-%m-%d"))
+        # Read the mod time
+        statinfo = None
+        try:
+            statinfo = os.stat(cache_file)
+        except:
+            pass
+
+        if statinfo is None or datetime.fromtimestamp(statinfo.st_mtime, tz=timezone.utc) + timedelta(minutes=self.cache_time) < datetime.now(tz=timezone.utc):
+            # Do the request and save it
+            req = urllib.request.Request(self.baseurl + "/1.0/infoVoy/rechercherListeCirculations?numero=%d&dateCirculation=%s&codeZoneArret&typeHoraire=TEMPS_REEL" % (numero, date.strftime("%Y-%m-%d")), headers={'Authorization': "Basic " + self.auth})
+            try:
+                with urllib.request.urlopen(req) as f:
+                    with open(cache_file, 'wb') as fd:
+                        fd.write(f.read())
+            except ConnectionResetError:
+                pass
+
+        # Retrieve cached data
+        res = {}
+        with open(cache_file) as f:
+            res = json.load(f)
+
+        return res["reponseRechercherListeCirculations"]["response"]["listeResultats"]["resultat"][0]["donnees"]["listeCirculations"]["circulation"][0]
+
+    def get_weather(self, region):
+        cache_file = self._cached_file % (0, region)
+        # Read the mod time
+        statinfo = None
+        try:
+            statinfo = os.stat(cache_file)
+        except:
+            pass
+
+        if statinfo is None or datetime.fromtimestamp(statinfo.st_mtime, tz=timezone.utc) + timedelta(minutes=self.cache_time) < datetime.now(tz=timezone.utc):
+            # Do the request and save it
+            req = urllib.request.Request(self.baseurl + "/edito/bandeaux?region=%s" % (region), headers={'Authorization': "Basic " + self.auth})
+            try:
+                with urllib.request.urlopen(req) as f:
+                    with open(cache_file, 'wb') as fd:
+                        fd.write(f.read())
+            except ConnectionResetError:
+                pass
+
+        # Retrieve cached data
+        res = {}
+        with open(cache_file) as f:
+            res = json.load(f)
+
+        return res
+
+
+class SNCFWeatherModule:
+
+    def __init__(self):
+        pass
+
+    def gen_alerts(self, region):
+        alerts = []
+
+        weather = SNCFAPI().get_weather(region)
+        for alert in weather:
+            if alert["type"] != "perturbation":
+                continue
+
+            yield {
+                "description": re.sub(SNCFAPI.CLEANR, '\n', alert["content"]).strip(),
+                "icon": SNCFAPI.get_icon,
+            }