Add support for checking limits and do some various improvements

Author: kkochanski

cli/arguments_builders/default_cli_arguments.py

@@ -13,33 +13,53 @@ def __init__(self, parser: ArgumentParser):
  def add_sha256_arg(self, help='Sample sha256 hash'):
  self.parser.add_argument('sha256', type=str, help=help)
 
+ return self
+
  def add_hash_arg(self, help='Md5, sha1 or sha256 hash'):
  self.parser.add_argument('hash', type=str, help=help)
 
+ return self
+
  def add_scan_type_arg(self):
  self.parser.add_argument('scan-type', type=str, help='Type of scan (please use \'{}\' action to fetch all available scanners)'.format(ACTION_SCAN_STATE))
 
+ return self
+
  def add_url_arg(self, help):
  self.parser.add_argument('url', type=str, help=help)
 
+ return self
+
  def add_id_arg(self, help='Id in one of format: \'jobId\' or \'sha256:environmentId\''):
  self.parser.add_argument('id', type=str, help=help)
 
+ return self
+
  def add_feed_days_arg(self):
  self.parser.add_argument('days', type=int, help='Number of days')
 
+ return self
+
  def add_key_uid_arg(self):
  self.parser.add_argument('uid', type=str, help='Any string to allow find this key later')
 
+ return self
+
  def add_file_with_hash_list_arg(self, allowed_hashes):
  self.parser.add_argument('hashes-file', type=argparse.FileType('r'), help='Path to file containing list of sample hashes separated by new line (allowed: {})'.format(', '.join(allowed_hashes)))
 
+ return self
+
  def add_file_with_ids_arg(self, allowed_ids):
  self.parser.add_argument('mixed-ids-file', type=argparse.FileType('r'), help='Path to file containing list of ids (allowed: {}'.format(', '.join(allowed_ids)))
 
+ return self
+
  def add_report_file_type_opt_arg(self):
  self.parser.add_argument('type', type=str, choices=['bin', 'json', 'pdf', 'crt', 'maec', 'stix', 'misp', 'misp-json', 'openioc', 'html', 'pcap', 'memory', 'xml'], help='Type of requested content')
 
+ return self
+
  def add_submit_files_arg(self):
  def validate_path(path):
  files = [path]
@@ -62,6 +82,8 @@ def validate_path(path):
 
  self.parser.add_argument('file', type=validate_path, help='File to submit (when directory given, all files from it will be submitted - non recursively)')
 
+ return self
+
  def add_file_output_path_opt(self):
  self.parser.add_argument('--output', '-o', type=str, default='output', help='File output path')
 
@@ -72,16 +94,26 @@ def add_env_id_arg(self, required: bool = False):
  else:
  self.parser.add_argument('environment-id', type=int, help=environment_id_help)
 
+ return self
+
  def add_report_file_type_opt(self):
  self.parser.add_argument('--type', '-t', type=str, choices=['bin', 'json', 'pdf', 'crt', 'maec', 'misp', 'misp-json', 'openioc', 'html', 'pcap', 'memory', 'xml'], help='File type to return')
 
+ return self
+
  def add_verbose_arg(self):
  self.parser.add_argument('--verbose', '-v', help="Run command in verbose mode", action='store_true')
 
+ return self
+
  def add_help_opt(self):
  self.parser.add_argument('--help', '-h', action='help', default=argparse.SUPPRESS, help='Show this help message and exit.')
 
+ return self
+
  def add_quiet_opt(self):
  self.parser.add_argument('--quiet', '-q', action='store_true', default=False, help='Suppress all prompts and warnings')
 
+ return self
+
 

cli/arguments_builders/demo_bulk_cli_arguments.py

@@ -6,11 +6,19 @@ class DemoBulkCliArguments(DefaultCliArguments):
  def add_report_demo_bulk_modify_hash_opt(self):
  self.parser.add_argument('--modify-hash', '-mh', action='store_true', default=False, help='When set, will add null byte at the end of sample file')
 
+ return self
+
  def add_report_demo_bulk_av_min_opt(self):
  self.parser.add_argument('--av-min', '-an', type=int, default=5, help='The minimum required AV detect')
 
+ return self
+
  def add_report_demo_bulk_av_max_opt(self):
  self.parser.add_argument('--av-max', '-ax', type=int, default=15, help='The maximum required AV detect')
 
+ return self
+
  def add_report_demo_bulk_look_back_size_opt(self):
  self.parser.add_argument('--look-back-size', '-lbs', type=int, default=400, help='Number of samples which will be fetched and filtered. Once you will get error message about problem with finding all samples, please increase that value')
+
+ return self

cli/arguments_builders/search_cli_arguments.py

@@ -7,21 +7,33 @@ class SearchCliArguments(DefaultCliArguments):
  def add_search_term_filename_opt(self):
  self.parser.add_argument('--filename', type=str, help='Filename e.g. invoice.exe')
 
+ return self
+
  def add_search_term_filetype_opt(self):
  self.parser.add_argument('--filetype', type=str, help='Filetype e.g. docx')
 
+ return self
+
  def add_search_term_filetype_desc_opt(self):
  self.parser.add_argument('--filetype-desc', type=str, help='Filetype description e.g. PE32 executable')
 
+ return self
+
  def add_search_term_env_id_opt(self):
  self.parser.add_argument('--env-id', type=int, help='Environment Id')
 
+ return self
+
  def add_search_term_country_opt(self):
  self.parser.add_argument('--country', type=str, help='Country (3 digit ISO) e.g. swe')
 
+ return self
+
  def add_search_term_verdict_opt(self):
  self.parser.add_argument('--verdict', type=int, help='Verdict', choices={1: 'whitelisted', 2: 'no verdict', 3: 'no specific threat', 4: 'suspicious', 5: 'malicious'})
 
+ return self
+
  def add_search_term_av_detect_opt(self):
  def type_av_detect(value):
  if value.find('-'):
@@ -38,41 +50,69 @@ def type_av_detect(value):
 
  self.parser.add_argument('--av-detect', type=type_av_detect, help='AV Multiscan range e.g. 50-70 (min 0, max 100)')
 
+ return self
+
  def add_search_term_vx_family_opt(self):
  self.parser.add_argument('--vx-family', type=str, help='AV Family Substring e.g. nemucod')
 
+ return self
+
  def add_search_term_tag_opt(self):
  self.parser.add_argument('--tag', type=str, help='Hashtag e.g. ransomware')
 
+ return self
+
  def add_search_term_date_from_opt(self):
  self.parser.add_argument('--date-to', type=str, help='Date from in format: ‘Y-m-d H:i’ e.g. 2018-09-28 15:30') # TODO - add some date validator here
 
+ return self
+
  def add_search_term_date_to_opt(self):
  self.parser.add_argument('--date-from', type=str, help='Date to in format: ‘Y-m-d H:i’ e.g. 2018-09-28 15:30')
 
+ return self
+
  def add_search_term_port_opt(self):
  self.parser.add_argument('--port', type=str, help='Port e.g. 8080')
 
+ return self
+
  def add_search_term_host_opt(self):
  self.parser.add_argument('--host', type=str, help='Host e.g. 192.168.0.1')
 
+ return self
+
  def add_search_term_domain_opt(self):
  self.parser.add_argument('--domain', type=str, help='Domain e.g. checkip.dyndns.org')
 
+ return self
+
  def add_search_term_url_opt(self):
  self.parser.add_argument('--url', type=str, help='HTTP Request Substring e.g. example')
 
+ return self
+
  def add_search_term_similar_to_opt(self):
  self.parser.add_argument('--similar-to', type=str, help='Similar Samples e.g. <sha266>')
 
+ return self
+
  def add_search_term_context_opt(self):
  self.parser.add_argument('--context', type=str, help='Sample Context e.g. <sha266>')
 
+ return self
+
  def add_search_term_imp_hash_opt(self):
  self.parser.add_argument('--imphash', type=str)
 
+ return self
+
  def add_search_term_ssdeep_opt(self):
  self.parser.add_argument('--ssdeep', type=str)
 
+ return self
+
  def add_search_term_authentihash_opt(self):
  self.parser.add_argument('--authentihash', type=str)
+
+ return self

cli/arguments_builders/submission_cli_arguments.py

@@ -7,59 +7,99 @@ class SubmissionCliArguments(DefaultCliArguments):
  def add_submission_submit_name_opt(self):
  self.parser.add_argument('--submit-name', '-sn', type=str, help='Optional \'submission name\' field that will be used for file type detection and analysis')
 
+ return self
+
  def add_submission_comment_opt(self):
  self.parser.add_argument('--comment', '-co', type=str, help='Add comment (e.g. #hashtag) to sample')
 
+ return self
+
  def add_submission_no_share_third_party_opt(self):
  self.parser.add_argument('--no-share-third-party', '-nstp', help='When set to \'1\', the sample is never shared with any third party', type=int, choices=[1, 0], default=1)
 
+ return self
+
  def add_submission_allow_community_access_opt(self):
  self.parser.add_argument('--allow-community-access', '-aca', choices=[1, 0], default=1, type=int, help='When set \'1\', the sample will be available for vetted users of the HA community or custom application server')
 
+ return self
+
  def add_submission_no_hash_lookup(self):
  self.parser.add_argument('--no-hash-lookup', '-nhl', choices=[1, 0], type=int)
 
+ return self
+
  def add_submission_action_script_opt(self):
  self.parser.add_argument('--action-script', '-ac', choices={1: 'default', 2: 'default_maxantievasion', 3: 'default_randomfiles', 4: 'default_randomtheme', 5: 'default_openie'}, type=int, help='Optional custom runtime action script')
 
+ return self
+
  def add_submission_hybrid_analysis_opt(self):
  self.parser.add_argument('--hybrid-analysis', '-ha', choices=[1, 0], type=int, help='When set to \'0\', no memory dumps or memory dump analysis will take place')
 
+ return self
+
  def add_submission_experimental_anti_evasion_opt(self):
  self.parser.add_argument('--experimental-anti-evasion', '-eae', choices=[1, 0], type=int, help='When set to \'1\', will set all experimental anti-evasion options of the Kernelmode Monitor')
 
+ return self
+
  def add_submission_script_logging_opt(self):
  self.parser.add_argument('--script-logging', '-sl', choices=[1, 0], type=int, help='When set to \'1\', will set the in-depth script logging engine of the Kernelmode Monitor')
 
+ return self
+
  def add_submission_input_sample_tampering_opt(self):
  self.parser.add_argument('--input-sample-tampering', '--ist', choices=[1, 0], type=int, help='When set to \'1\', will allow experimental anti-evasion options of the Kernelmode Monitor that tamper with the input sample')
 
+ return self
+
  def add_submission_tor_enabled_analysis_opt(self):
  self.parser.add_argument('--tor-enabled-analysis', '-tea', choices=[1, 0], type=int, help='When set to \'1\', will route the network traffic for the analysis via TOR (if properly configured on the server)')
 
+ return self
+
  def add_submission_offline_analysis_opt(self):
  self.parser.add_argument('--offline-analysis', '-oa', choices=[1, 0], type=int, help='When set to \'1\', will disable outbound network traffic for the guest VM (takes precedence over ‘tor-enabled-analysis’ if both are provided)')
 
+ return self
+
  def add_submission_email_opt(self):
  self.parser.add_argument('--email', '-e', type=str, help='Optional E-Mail address that may be associated with the submission for notification')
 
+ return self
+
  def add_submission_custom_date_time_opt(self):
  self.parser.add_argument('--custom-date-time', '-cdt', type=str, help='Optional custom date/time that can be set for the analysis system. Expected format: yyyy-MM-dd HH:mm')
 
+ return self
+
  def add_submission_custom_cmd_line_opt(self):
  self.parser.add_argument('--custom-cmd-line', '-ccl', type=str, help='Optional commandline that should be passed to the analysis file')
 
+ return self
+
  def add_submission_custom_run_time_opt(self):
  self.parser.add_argument('--custom-run-time', '-crt', type=int, help='Optional runtime duration (in seconds)',)
 
+ return self
+
  def add_submission_client_opt(self):
  self.parser.add_argument('--client', '-cl', type=str, help='Optional ‘client’ field (see ‘vxClients’)')
 
+ return self
+
  def add_submission_priority_opt(self):
  self.parser.add_argument('--priority', '-pr', type=ValuesInBetweenAction(), help='Optional priority value between 0 (default) and 100 (highest)')
 
+ return self
+
  def add_submission_document_password_opt(self):
  self.parser.add_argument('--document-password', '-dp', type=str, help='Optional document password that will be used to fill-in Adobe/Office password prompts')
 
+ return self
+
  def add_submission_environment_variable_opt(self):
  self.parser.add_argument('--environment-variable', '-ev', type=str, help='Optional system environment value. The value is provided in the format: name=value')
+
+ return self

cli/cli_msg_printer.py

@@ -38,11 +38,27 @@ def print_error_info(e):
 
  @staticmethod
  def print_usage_info(api_usage_limits, api_usage, is_api_limit_reached):
- print(Color.control('API Limits for used API Key'))
+ print(Color.control('API query limits for used API Key'))
  print('Webservice API usage limits: {}'.format(api_usage_limits))
  print('Current API usage: {}'.format(json.dumps(api_usage)))
  print('Is limit reached: {}'.format(Color.success('No') if is_api_limit_reached is False else Color.error('Yes')))
 
+ @staticmethod
+ def print_submission_limit_info(submission_limits_data):
+ if submission_limits_data:
+ print(Color.control('Submission limits for used API Key'))
+ print('Webservice API usage limits: {}'.format(submission_limits_data['total']['quota']))
+ print('Current API usage: {}'.format(json.dumps(submission_limits_data['total']['used'])))
+ print('Is limit reached: {}'.format(Color.success('No') if submission_limits_data['total']['quota_reached'] is False else Color.error('Yes')))
+
+ @staticmethod
+ def print_quick_scan_limit_info(quick_scan_limits_data):
+ if quick_scan_limits_data:
+ print(Color.control('Quick scan submission limits for used API Key'))
+ print('Webservice API usage limits: {}'.format(quick_scan_limits_data['total']['quota']))
+ print('Current API usage: {}'.format(json.dumps(quick_scan_limits_data['total']['used'])))
+ print('Is limit reached: {}'.format(Color.success('No') if quick_scan_limits_data['total']['quota_reached'] is False else Color.error('Yes')))
+
  @staticmethod
  def print_api_key_info(current_key_json):
  print(Color.control('Used API Key'))

constants.py

@@ -50,6 +50,8 @@
 ACTION_SCAN_URL_FOR_ANALYSIS = 'scan_url_for_analysis'
 ACTION_SCAN_URL_TO_FILE = 'scan_url_to_file'
 
+ACTION_CHECK_LIMITS = 'check_limits'
+
 MINIMAL_SUPPORTED_INSTANCE_VERSION = '8.20'
 
 ACTION_WITH_MULTIPLE_CALL_SUPPORT = [ACTION_SUBMIT_FILE, ACTION_SCAN_FILE]

tests/base_test.py

@@ -32,7 +32,7 @@ def do_run(*args):
 
  def see_headers(self):
  assert 'Running \'VxWebService Python API Connector\'' in self.output
- assert 'API Limits for used API Key' in self.output
+ assert 'API query limits for used API Key' in self.output
  assert 'Request was sent at' in self.output
  assert 'Received response at' in self.output
  assert 'Showing response' in self.output

tests/conftest.py

@@ -3,6 +3,6 @@
 
 os.environ['APP_ENV'] = 'test'
 os.environ['TEST_CONFIG'] = json.dumps({
- 'api_key': 'test_me_please',
+ 'api_key': '6b86b273ff34fce19d6b804eff5a3f5747ada4eaa22f1d49c01e52ddb7875b4b',
  'server': 'mock://my-webservice-instance'
 })