1import requests
2import json
3import ssl
4
5from .exception import APIError
6
7
8class Stream:
9
10    base_url = 'https://stream.shodan.io'
11
12    def __init__(self, api_key, proxies=None):
13        self.api_key = api_key
14        self.proxies = proxies
15
16    def _create_stream(self, name, timeout=None):
17        params = {
18            'key': self.api_key,
19        }
20        stream_url = self.base_url + name
21
22        # The user doesn't want to use a timeout
23        # If the timeout is specified as 0 then we also don't want to have a timeout
24        if (timeout and timeout <= 0) or (timeout == 0):
25            timeout = None
26
27        # If the user requested a timeout then we need to disable heartbeat messages
28        # which are intended to keep stream connections alive even if there isn't any data
29        # flowing through.
30        if timeout:
31            params['heartbeat'] = False
32
33        try:
34            while True:
35                req = requests.get(stream_url, params=params, stream=True, timeout=timeout,
36                                   proxies=self.proxies)
37
38                # Status code 524 is special to Cloudflare
39                # It means that no data was sent from the streaming servers which caused Cloudflare
40                # to terminate the connection.
41                #
42                # We only want to exit if there was a timeout specified or the HTTP status code is
43                # not specific to Cloudflare.
44                if req.status_code != 524 or timeout >= 0:
45                    break
46        except Exception:
47            raise APIError('Unable to contact the Shodan Streaming API')
48
49        if req.status_code != 200:
50            try:
51                data = json.loads(req.text)
52                raise APIError(data['error'])
53            except APIError:
54                raise
55            except Exception:
56                pass
57            raise APIError('Invalid API key or you do not have access to the Streaming API')
58        if req.encoding is None:
59            req.encoding = 'utf-8'
60        return req
61
62    def _iter_stream(self, stream, raw):
63        for line in stream.iter_lines(decode_unicode=True):
64            # The Streaming API sends out heartbeat messages that are newlines
65            # We want to ignore those messages since they don't contain any data
66            if line:
67                if raw:
68                    yield line
69                else:
70                    yield json.loads(line)
71
72    def alert(self, aid=None, timeout=None, raw=False):
73        if aid:
74            stream = self._create_stream('/shodan/alert/%s' % aid, timeout=timeout)
75        else:
76            stream = self._create_stream('/shodan/alert', timeout=timeout)
77
78        try:
79            for line in self._iter_stream(stream, raw):
80                yield line
81        except requests.exceptions.ConnectionError:
82            raise APIError('Stream timed out')
83        except ssl.SSLError:
84            raise APIError('Stream timed out')
85
86    def asn(self, asn, raw=False, timeout=None):
87        """
88        A filtered version of the "banners" stream to only return banners that match the ASNs of interest.
89
90        :param asn: A list of ASN to return banner data on.
91        :type asn: string[]
92        """
93        stream = self._create_stream('/shodan/asn/%s' % ','.join(asn), timeout=timeout)
94        for line in self._iter_stream(stream, raw):
95            yield line
96
97    def banners(self, raw=False, timeout=None):
98        """A real-time feed of the data that Shodan is currently collecting. Note that this is only available to
99        API subscription plans and for those it only returns a fraction of the data.
100        """
101        stream = self._create_stream('/shodan/banners', timeout=timeout)
102        for line in self._iter_stream(stream, raw):
103            yield line
104
105    def countries(self, countries, raw=False, timeout=None):
106        """
107        A filtered version of the "banners" stream to only return banners that match the countries of interest.
108
109        :param countries: A list of countries to return banner data on.
110        :type countries: string[]
111        """
112        stream = self._create_stream('/shodan/countries/%s' % ','.join(countries), timeout=timeout)
113        for line in self._iter_stream(stream, raw):
114            yield line
115
116    def ports(self, ports, raw=False, timeout=None):
117        """
118        A filtered version of the "banners" stream to only return banners that match the ports of interest.
119
120        :param ports: A list of ports to return banner data on.
121        :type ports: int[]
122        """
123        stream = self._create_stream('/shodan/ports/%s' % ','.join([str(port) for port in ports]), timeout=timeout)
124        for line in self._iter_stream(stream, raw):
125            yield line
126
127    def tags(self, tags, raw=False, timeout=None):
128        """
129        A filtered version of the "banners" stream to only return banners that match the tags of interest.
130
131        :param tags: A list of tags to return banner data on.
132        :type tags: string[]
133        """
134        stream = self._create_stream('/shodan/tags/%s' % ','.join(tags), timeout=timeout)
135        for line in self._iter_stream(stream, raw):
136            yield line
137
138    def vulns(self, vulns, raw=False, timeout=None):
139        """
140        A filtered version of the "banners" stream to only return banners that match the vulnerabilities of interest.
141
142        :param vulns: A list of vulns to return banner data on.
143        :type vulns: string[]
144        """
145        stream = self._create_stream('/shodan/vulns/%s' % ','.join(vulns), timeout=timeout)
146        for line in self._iter_stream(stream, raw):
147            yield line
148