def queue_bulk_threats(self, atom_list, payload):
hashkey_created = []
bulk_in_flight = [] # bulk task uuid unchecked
for batch in split_list(atom_list, self._batch_size()):
if len(bulk_in_flight) >= self.OCD_DTL_MAX_BULK_THREATS_IN_FLIGHT:
bulk_threat_task_uuid = bulk_in_flight.pop(0)
hashkey_created += self.check_bulk_threats_added(bulk_threat_task_uuid)
payload['atom_values'] = '\n'.join(batch) # Raw csv expected
response = self.datalake_requests(self.url, 'post', self._post_headers(), payload)
task_uid = response.get('task_uuid')
if task_uid:
bulk_in_flight.append(response['task_uuid'])
else:
logger.warning(f'batch of threats from {batch[0]} to {batch[-1]} failed to be created')
# Finish to check the other bulk tasks
for bulk_threat_task_uuid in bulk_in_flight:
hashkey_created += self.check_bulk_threats_added(bulk_threat_task_uuid)
nb_threats = len(hashkey_created)
if nb_threats > 0:
ok_sign = '\x1b[0;30;42m' + ' OK ' + '\x1b[0m'
logger.info(f'Created {nb_threats} threats'.ljust(self.terminal_size - 6, ' ') + ok_sign)
else:
ko_sign = '\x1b[0;30;41m' + ' KO ' + '\x1b[0m'
logger.info(f'Failed to create any threats'.ljust(self.terminal_size - 6, ' ') + ko_sign)
return set(hashkey_created)
def test_split_list():
list_to_split = [1, 2, 3, 4, 5]
generator = split_list(list_to_split, 3)
head = next(generator)
tail = next(generator)
assert head == [1, 2, 3]
assert tail == [4, 5]
def main(override_args=None):
"""Method to start the script"""
# Load initial args
parser = BaseScripts.start(
'Edit scores of a specified list of ids (hashkeys)')
parser.add_argument(
'hashkeys',
help='hashkeys of the threat to edit score.',
nargs='*',
)
parser.add_argument(
'-i',
'--input_file',
help='hashkey txt file, with one hashkey by line.',
)
parser.add_argument(
'-t',
'--threat_types',
nargs='+',
help=
'choose specific threat types and their score, like: ddos 50 scam 15',
default=[],
action='append',
)
parser.add_argument(
'-w',
'--whitelist',
help=
'Whitelist the input, equivalent to setting all threat types at 0.',
action='store_true',
)
parser.add_argument(
'--permanent',
help=
'''Permanent: all values will override any values provided by both newer and
older IOCs. Newer IOCs with override_type permanent can still override old permanent changes.
temporary: all values should override any values provided by older IOCs,
but not newer ones.''',
action='store_true',
)
parser.add_argument(
'--lock',
help=
'sets override_type to lock. Scores won\'t be updated by the algorithm for three months. Default is '
'temporary',
action='store_true',
)
if override_args:
args = parser.parse_args(override_args)
else:
args = parser.parse_args()
logger.debug(f'START: edit_score.py')
if not args.hashkeys and not args.input_file:
parser.error("either a hashkey or an input_file is required")
if args.permanent and args.lock:
parser.error("Only one override type is authorized")
if args.permanent:
override_type = OverrideType.PERMANENT
elif args.lock:
override_type = OverrideType.LOCK
else:
override_type = OverrideType.TEMPORARY
if args.whitelist:
parsed_threat_type = get_whitelist_threat_types()
else:
args.threat_types = flatten_list(args.threat_types)
if not args.threat_types or len(args.threat_types) % 2 != 0:
parser.error(
"threat_types invalid ! should be like: ddos 50 scam 15")
parsed_threat_type = parse_threat_types(args.threat_types)
# removing duplicates while preserving order
hashkeys = args.hashkeys
if args.input_file:
retrieve_hashkeys_from_file(args.input_file, hashkeys)
if not hashkeys:
raise parser.error('No hashkey found in the input file.')
hashkeys_chunks = list(
split_list(
list(OrderedDict.fromkeys(hashkeys)) if hashkeys else [], 100))
dtl = Datalake(env=args.env, log_level=args.loglevel)
response_list = []
for index, hashkeys in enumerate(hashkeys_chunks):
try:
dtl.Threats.edit_score_by_hashkeys(hashkeys, parsed_threat_type,
override_type)
except ValueError as e:
logger.warning(
f'\x1b[6;30;41mBATCH {str(index+1)}/{len(list(hashkeys_chunks))}: FAILED\x1b[0m'
)
for hashkey in hashkeys:
response_list.append(hashkey + ': FAILED')
logger.warning(f'\x1b[6;30;41m{hashkey} : FAILED\x1b[0m')
logger.warning(e)
else:
logger.info(
f'\x1b[6;30;42mBATCH {str(index+1)}/{len(list(hashkeys_chunks))}: OK\x1b[0m'
)
for hashkey in hashkeys:
response_list.append(hashkey + ': OK')
if args.output:
save_output(args.output, response_list)
logger.info(f'Results saved in {args.output}\n')
logger.debug(f'END: edit_score.py')