From b448c516412ebc9fda56c7043ab290e3fedc44b8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mikk=20Margus=20M=C3=B6ll?= <mikk@cert.ee>
Date: Wed, 10 Apr 2024 14:07:04 +0300
Subject: [PATCH] ENH: add alert parameter to Shodan stream collector

---
 CHANGELOG.md                                  |  2 ++
 docs/user/bots.md                             |  4 ++++
 .../collectors/shodan/collector_stream.py     | 23 +++++++++++++++----
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 810ad8deb..08b67499d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,6 +19,8 @@
 
 ### Bots
 #### Collectors
+- `intelmq.bots.collectors.shodan.collector_stream` (PR#2492 by Mikk Margus Möll):
+  - Add `alert` parameter to Shodan stream collector to allow fetching streams by configured alert ID
 
 #### Parsers
 - `intelmq.bots.parsers.shadowserver._config`:
diff --git a/docs/user/bots.md b/docs/user/bots.md
index 6d44d7cf8..2c83de069 100644
--- a/docs/user/bots.md
+++ b/docs/user/bots.md
@@ -994,6 +994,10 @@ Only the proxy is used (requires `shodan-python > 1.8.1`). Certificate is always
 
 () A list of countries to query for. If it is a string, it will be spit by `,`.
 
+**`alert`**
+
+() Alert ID from monitor.shodan.io.
+
 If the stream is interrupted, the connection will be aborted using the timeout parameter. No error will be logged if the
 number of consecutive connection fails does not reach the parameter
 `error_max_retries`. Instead of errors, an INFO message is logged. This is a measurement against too frequent ERROR
diff --git a/intelmq/bots/collectors/shodan/collector_stream.py b/intelmq/bots/collectors/shodan/collector_stream.py
index aeed4ff26..9b17cac60 100644
--- a/intelmq/bots/collectors/shodan/collector_stream.py
+++ b/intelmq/bots/collectors/shodan/collector_stream.py
@@ -8,15 +8,15 @@
 * api_key: The API key
 
 Selectors:
-The only possible selector is currently the country:
 * countries: A list of strings or a comma separated list with country codes
+* alert: An alert ID from monitor.shodan.io
 """
 import pkg_resources
 from http.client import IncompleteRead
 from urllib3.exceptions import ProtocolError, ReadTimeoutError
 
 from requests.exceptions import ChunkedEncodingError, ConnectionError
-from typing import List
+from typing import List, Optional
 
 from intelmq.lib.bot import CollectorBot
 
@@ -31,6 +31,7 @@ class ShodanStreamCollectorBot(CollectorBot):
     "Collect the Shodan stream from the Shodan API"
     api_key: str = "<INSERT your API key>"
     countries: List[str] = []
+    alert: Optional[str] = None
 
     def init(self):
         if shodan is None:
@@ -46,14 +47,28 @@ def init(self):
             self.api = shodan.Shodan(self.api_key,
                                      proxies=self.proxy)
         if isinstance(self.countries, str):
+            if self.countries and self.alert:
+                raise ValueError('Both alert and country filters specified. Please use only one selector.')
             self.countries = self.countries.split(',')
 
         self.__error_count = 0
 
     def process(self):
         try:
-            for line in self.api.stream.countries(timeout=self.http_timeout_sec, raw=True,
-                                                  countries=self.countries):
+            if self.alert:
+                stream = self.api.stream.alert(
+                    aid=self.alert,
+                    timeout=self.http_timeout_sec,
+                    raw=True,
+                )
+            else:
+                stream = self.api.stream.countries(
+                    countries=self.countries,
+                    timeout=self.http_timeout_sec,
+                    raw=True,
+                )
+
+            for line in stream:
                 report = self.new_report()
                 report.add('raw', line)
                 self.send_message(report)