async def stop_sds_collection(self, sds_p):
sds_p.disable = True
try:
self.sd_file.close()
cnt = 0
with gzip.GzipFile(self.sd_file_path, 'r') as f:
for line in tqdm(f, desc='Uploading security descriptors to DB', total=self.spn_finish_ctr):
sd = JackDawSD.from_json(line.strip())
self.session.add(sd)
cnt += 1
if cnt % 10000 == 0:
self.session.commit()
self.session.commit()
os.remove(self.sd_file_path)
except Exception as e:
logger.exception('Error while uploading sds from file to DB')
async def store_file_data(self):
try:
self.progress_last_updated = datetime.datetime.utcnow()
if self.progress_queue is not None:
msg = GathererProgress()
msg.type = GathererProgressType.SDUPLOAD
msg.msg_type = MSGTYPE.STARTED
msg.adid = self.ad_id
msg.domain_name = self.domain_name
await self.progress_queue.put(msg)
if self.show_progress is True:
self.sd_upload_pbar = tqdm(desc='uploading SD to DB',
total=self.total_targets)
if self.sd_file is not None:
self.sd_file.close()
cnt = 0
buffer = []
last_stat_cnt = 0
with gzip.GzipFile(self.sd_file_path, 'r') as f:
for line in f:
buffer.append(JackDawSD.from_json(line.strip()))
await asyncio.sleep(0)
cnt += 1
if cnt % 100 == 0:
self.session.bulk_save_objects(buffer)
buffer = []
if self.show_progress is True:
self.sd_upload_pbar.update()
if self.progress_queue is not None and cnt % self.progress_step_size == 0:
last_stat_cnt += self.progress_step_size
now = datetime.datetime.utcnow()
td = (now -
self.progress_last_updated).total_seconds()
self.progress_last_updated = now
msg = GathererProgress()
msg.type = GathererProgressType.SDUPLOAD
msg.msg_type = MSGTYPE.PROGRESS
msg.adid = self.ad_id
msg.domain_name = self.domain_name
msg.total = self.total_targets
msg.total_finished = cnt
if td > 0:
msg.speed = str(self.progress_step_size // td)
msg.step_size = self.progress_step_size
await self.progress_queue.put(msg)
await asyncio.sleep(0)
if len(buffer) > 0:
self.session.bulk_save_objects(buffer)
buffer = []
self.session.commit()
if self.progress_queue is not None:
now = datetime.datetime.utcnow()
td = (now - self.progress_last_updated).total_seconds()
self.progress_last_updated = now
msg = GathererProgress()
msg.type = GathererProgressType.SDUPLOAD
msg.msg_type = MSGTYPE.PROGRESS
msg.adid = self.ad_id
msg.domain_name = self.domain_name
msg.total = self.total_targets
msg.total_finished = cnt
if td > 0:
msg.speed = str(
(self.total_targets - last_stat_cnt) // td)
msg.step_size = self.total_targets - last_stat_cnt
await self.progress_queue.put(msg)
await asyncio.sleep(0)
except Exception as e:
logger.exception('Error while uploading sds from file to DB')
if self.progress_queue is not None:
msg = GathererProgress()
msg.type = GathererProgressType.SDUPLOAD
msg.msg_type = MSGTYPE.ERROR
msg.adid = self.ad_id
msg.domain_name = self.domain_name
msg.error = e
await self.progress_queue.put(msg)
finally:
try:
os.remove(self.sd_file_path)
except:
pass
if self.show_progress is True and self.sd_upload_pbar is not None:
self.sd_upload_pbar.refresh()
self.sd_upload_pbar.disable = True
if self.progress_queue is not None:
msg = GathererProgress()
msg.type = GathererProgressType.SDUPLOAD
msg.msg_type = MSGTYPE.FINISHED
msg.adid = self.ad_id
msg.domain_name = self.domain_name
await self.progress_queue.put(msg)