[FIX] Incomplete Shodan-API error message.

Xyntax · Mar 25, 2017 · a3a67b6 · a3a67b6
1 parent 60c290c
commit a3a67b6
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 31 deletions.
diff --git a/lib/api/shodan/pack.py b/lib/api/shodan/pack.py
@@ -10,35 +10,55 @@
 from lib.utils.config import ConfigFileParser
 
 
-def _readKey():
-    msg = 'Trying to auth with credentials in config file: %s.' % paths.CONFIG_PATH
-    logger.info(msg)
-    try:
-        key = ConfigFileParser().ShodanApikey()
-    except:
-        key = ''
-    return key
+class ShodanBase:
+    def __init__(self, query, limit, offset):
+        self.query = query
+        self.limit = limit
+        self.offset = offset
+        self.api_key = None
+        self.result = None
 
+    def login(self):
+        msg = 'Trying to login with credentials in config file: %s.' % paths.CONFIG_PATH
+        logger.info(msg)
+        self.api_key = ConfigFileParser().ShodanApikey()
 
-def ShodanSearch(query, limit, offset=0):
-    key = _readKey()
-    try:
-        api = shodan.Shodan(key)
-        result = api.search(query=query, offset=offset, limit=limit)
-    except APIError:
-        msg = 'Automatic authorization failed.'
-        logger.warning(msg)
-        key = raw_input('Input API-KEY >')
+        if not self.api_key:
+            msg = 'Automatic authorization failed.'
+            logger.warning(msg)
+            msg = 'Please input your Shodan API Key (https://account.shodan.io/).'
+            logger.info(msg)
+            self.api_key = raw_input('API KEY > ').strip()
+
+    def account_info(self):
+        try:
+            api = shodan.Shodan(self.api_key)
+            account_info = api.info()
+            msg = "Available Shodan query credits: %d" % account_info['query_credits']
+            logger.info(msg)
+        except APIError, e:
+            sys.exit(logger.error(e))
+        return True
+
+    def api_query(self):
         try:
-            api = shodan.Shodan(key)
-            result = api.search(query=query, offset=offset, limit=limit)
-        except APIError:
-            msg = 'Invalid Shodan API key.'
-            sys.exit(logger.error(msg))
-    if 'matches' in result:
-        anslist = []
-        for match in result['matches']:
-            anslist.append(match['ip_str'] + ':' + str(match['port']))
-        return anslist
-    else:
-        return []
+            api = shodan.Shodan(self.api_key)
+            result = api.search(query=self.query, offset=self.offset, limit=self.limit)
+        except APIError, e:
+            sys.exit(logger.error(e))
+
+        if 'matches' in result:
+            anslist = []
+            for match in result['matches']:
+                anslist.append(match['ip_str'] + ':' + str(match['port']))
+            self.result = anslist
+        else:
+            self.result = []
+
+
+def ShodanSearch(query, limit, offset=0):
+    s = ShodanBase(query, limit, offset)
+    s.login()
+    s.account_info()
+    s.api_query()
+    return s.result
diff --git a/lib/api/zoomeye/base.py b/lib/api/zoomeye/base.py
@@ -39,8 +39,8 @@ def auto_login(self):
     def manual_login(self):
         msg = 'Please input your ZoomEye Email and Password below.'
         logger.info(msg)
-        self.username = raw_input('ZoomEye Username(Email): ')
-        self.password = getpass.getpass(prompt='ZoomEye Password: ')
+        self.username = raw_input('ZoomEye Username(Email): ').strip()
+        self.password = getpass.getpass(prompt='ZoomEye Password: ').strip()
         if not self.get_token():
             msg = 'Invalid ZoomEye username or password.'
             sys.exit(logger.error(msg))