chg: [FediVuln-Publish] It is now possible to stream from Valkey inst…

…ead of streaming for the HTTP event-stream.

fedivuln/conf_sample.py

@@ -1,17 +1,15 @@
+# ### Mastodon
+
+# Connection to Mastodon
 api_base_url = "https://social.circl.lu"
 scopes = ["read", "write", "follow", "push"]
 app_name = "Vulnerability-Lookup"
 mastodon_clientcred = "mastodon_clientcred.secret"
 mastodon_usercred = "mastodon_usercred.secret"
 
-# Optional in case you need to publish status with a different account
-mastodon_clientcred_push = "mastodon_clientcred_push.secret"
-mastodon_usercred_push = "mastodon_usercred_push.secret"
-
-
-vulnerability_lookup_base_url = "https://vulnerability.circl.lu/"
-vulnerability_auth_token = ""
-
+# Optional in case you need to publish status with a different account than the one previously defined:
+mastodon_clientcred_push = ""
+mastodon_usercred_push = ""
 
 # Templates used when publishing status
 templates = {
@@ -22,3 +20,14 @@
     "bundle": "A new bundle, <BUNDLETITLE>, has been published "
     "on Vulnerability-Lookup:\n<LINK>\n\n#VulnerabilityLookup #Vulnerability #Cybersecurity #bot",
 }
+
+
+# ### Event stream
+
+# Choice 1: Stream from the authenticated HTTP event stream of Vulnerability-Lookup (default):
+vulnerability_lookup_base_url = "https://vulnerability.circl.lu/"
+vulnerability_auth_token = ""
+
+# Choice 2: Stream from the Valkey Pub/Sub streaming service (--valkey option):
+valkey_host = "127.0.0.1"
+valkey_port = 10002

fedivuln/config.py

@@ -45,6 +45,9 @@ def load_config(path):
     # Optional second Mastodon account to publish status
     mastodon_clientcred_push = conf.mastodon_clientcred_push
     mastodon_usercred_push = conf.mastodon_usercred_push
+
+    valkey_host = conf.valkey_host
+    valkey_port = conf.valkey_port
 except AttributeError:
     mastodon_clientcred_push = ""
     mastodon_usercred_push = ""

fedivuln/publish.py

@@ -3,6 +3,7 @@
 from urllib.parse import urljoin
 
 import requests
+import valkey
 from mastodon import Mastodon
 
 from fedivuln import config
@@ -90,9 +91,41 @@ def listen_to_http_event_stream(url, headers=None, params=None, topic="comment")
         print(f"Unexpected error: {e}")
 
 
+def listen_to_valkey_stream(topic="comment"):
+    """Stream data from the Valkey Pub/Sub service."""
+    valkey_client = valkey.Valkey(
+        host=config.valkey_host,
+        port=config.valkey_port,
+        decode_responses=True,
+    ).pubsub()
+    try:
+        valkey_client.subscribe(topic)
+    except valkey.exceptions.ConnectionError:
+        return
+    try:
+        while True:
+            message = valkey_client.get_message(timeout=10)  # Timeout for listener
+            if message and message["type"] == "message":
+                # Send entire JSON object as a single `data:` line
+                json_message = json.dumps(message["data"])  # Ensure single-line JSON
+                yield f"{json_message}"
+    except GeneratorExit:
+        valkey_client.unsubscribe(topic)
+    except valkey.exceptions.ConnectionError:
+        return
+    finally:
+        valkey_client.close()
+
+
 def main():
     """Parsing of arguments."""
     parser = argparse.ArgumentParser(prog="FediVuln-Publish")
+    parser.add_argument(
+        "--valkey",
+        dest="valkey",
+        action="store_true",
+        help="Stream from Valkey instead of streaming for the HTTP event-stream.",
+    )
     parser.add_argument(
         "-t",
         "--topic",
@@ -104,10 +137,15 @@ def main():
 
     arguments = parser.parse_args()
 
-    combined = urljoin(config.vulnerability_lookup_base_url, "pubsub/subscribe/")
-    full_url = urljoin(combined, arguments.topic)
-    headers = {"X-API-KEY": config.vulnerability_auth_token}
-    listen_to_http_event_stream(full_url, headers=headers, topic=arguments.topic)
+    if arguments.valkey:
+        for elem in listen_to_valkey_stream(topic=arguments.topic):
+            event_data = json.loads(elem)
+            publish(create_status_content(event_data, arguments.topic))
+    else:
+        combined = urljoin(config.vulnerability_lookup_base_url, "pubsub/subscribe/")
+        full_url = urljoin(combined, arguments.topic)
+        headers = {"X-API-KEY": config.vulnerability_auth_token}
+        listen_to_http_event_stream(full_url, headers=headers, topic=arguments.topic)
 
 
 if __name__ == "__main__":

fedivuln/stream.py

@@ -114,33 +114,6 @@ def push_sighting_to_vulnerability_lookup(status_uri, vulnerability_ids):
             )
 
 
-# def push_status_to_vulnerability_lookup(status, vulnerability_ids):
-#     """Push the status to the Vulnerability Lookup instance."""
-#     print("Pushing status to Vulnerability Lookup…")
-#     headers_json = {
-#         "Content-Type": "application/json",
-#         "accept": "application/json",
-#         "X-API-KEY": f"{config.vulnerability_auth_token}",
-#     }
-#     # status = convert_datetime(status)
-#     json_payload = {"status": status, "vulnerability_ids": vulnerability_ids}
-#     json_string = json.dumps(json_payload, cls=DateTimeEncoder)
-#     try:
-#         r = requests.post(
-#             urllib.parse.urljoin(config.vulnerability_lookup_base_url, "status/"),
-#             data=json_string,
-#             headers=headers_json,
-#         )
-#         if r.status_code not in (200, 201):
-#             print(
-#                 f"Error when sending POST request to the Vulnerability Lookup server: {r.reason}"
-#             )
-#     except requests.exceptions.ConnectionError as e:
-#         print(
-#             f"Error when sending POST request to the Vulnerability Lookup server:\n{e}"
-#         )
-
-
 def main():
     parser = argparse.ArgumentParser(
         prog="FediVuln-Stream", description="Allows access to the streaming API."

pyproject.toml

@@ -41,6 +41,7 @@ FediVuln-Publish = "fedivuln.publish:main"
 python = "^3.10"
 mastodon-py = "^1.8.1"
 pyvulnerabilitylookup = "^2.1.0"
+valkey = "^6.0.2"
 
 [tool.poetry.group.dev.dependencies]
 mypy = "^1.13.0"