def do_source(self, arg): """Add source to given artifact or most recently added artifact if not specified Usage: source # adds to last created artifact source <artifact name|session id> # adds to specific artifact """ if arg == '': last = self.session.receive('artifacts') _type = detect_type(last) else: _type = detect_type(arg) is_key, value = lookup_key(self.session, arg) if is_key and value is None: error('Unable to find artifact key in session (%s)' % arg) return elif is_key and value is not None: arg = value else: pass if self.db.exists(_type, {'name': last}): self.db.update_one(_type, {'name': last}, {'source': arg}) success('Added source to artifact entry (%s: %s)' % (last, arg)) else: warning('Failed to find last artifact in MongoDB. Run "new <artifact name>" before using the source command')
def do_open(self, arg): """Load text file list of artifacts Command will detect each line items artifact type, create the artifact, and add it to the current session if there is one. Usage: open <path/to/file.txt> """ if not os.path.exists(arg): warning('Cannot find file on disk (%s)' % arg) return artifacts = read_file(arg, True) for artifact in artifacts: new_artifact = create_artifact(artifact) if not self.db.exists(new_artifact.type, {'name': new_artifact.name}): doc_id = self.db.insert_one(new_artifact.type, new_artifact) if doc_id is not None: success('Created new artifact (%s - %s)' % (artifact.name, artifact.type)) if self.session is None: self.session = RedisCache(config) self.session.set(1, arg) success('Opened new session') print('Artifact ID: 1') else: count = 0 for key in self.session.db.scan_iter(): count += 1 _id = count + 1 self.session.set(_id, arg) print('Artifact ID: %s' % _id) success('Finished loading artifact list')
def do_report(self, arg): """Save artifact report as JSON file Usage: report <artifact name> report <session id>""" is_key, value = lookup_key(self.session, arg) if is_key and value is None: error('Unable to find artifact key in session (%s)' % arg) return elif is_key and value is not None: arg = value else: pass _type = detect_type(arg) result = self.db.find(_type, {'name': arg}, one=True) if len(result) == 0: warning('No entry found for artifact (%s)' % arg) else: report = storage.JSON(data=result, file_path=output_dir) report.save() if os.path.exists(report.file_path): success('Saved artifact report (%s)' % report.file_path) else: error('Failed to properly save report')
def do_wipe(self, arg): """Clear currently active artifacts """ if self.session is not None: info('Clearing active artifacts from cache ...') self.session.flush() success('Artifact cache cleared') else: warning('No active session; start a new session by running the "session" command')
def do_quit(self, _): """Exit Omnibus shell.""" self._should_quit = True if self.session is not None: running('Clearing artifact cache ...') self.session.flush() warning('Closing Omnibus shell ...') return self._STOP_AND_EXIT
def do_ls(self, arg): """View current sessions artifacts""" if self.session is None: warning('No active session') return count = 0 keys = self.session.db.scan_iter() for key in keys: value = self.session.get(key) print('[%s] %s' % (key, value)) count += 1 info('Active Artifacts: %d' % count)
def do_rm(self, arg): """Remove artifact from session by ID Usage: rm <session id>""" try: arg = int(arg) except: error('Artifact ID must be an integer') return if self.session is not None: if self.session.exists(arg): self.session.delete(arg) success('Removed artifact from cache (%s)' % arg) else: warning('Unable to find artifact by ID (%s)' % arg) else: warning('No active session; start a new session by running the "session" command')