diff --git a/lib/api/shodan/pack.py b/lib/api/shodan/pack.py index b14045d..6fad7c6 100644 --- 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 index ef0ff13..a521fcd 100644 --- 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))