class DB: def __init__(self): self.badges = Badges() self.local_storage = LocalStorage() def disconnect_payloads_database(self, name): if self.local_storage.get("connected_payloads_databases"): if name in self.local_storage.get("connected_payloads_databases"): self.local_storage.delete_element( "connected_payloads_databases", name) self.local_storage.delete_element("payloads", name) return self.badges.output_error("No such payloads database connected!") def disconnect_modules_database(self, name): if self.local_storage.get("connected_modules_databases"): if name in self.local_storage.get("connected_modules_databases"): self.local_storage.delete_element( "connected_modules_databases", name) self.local_storage.delete_element("modules", name) return self.badges.output_error("No such modules database connected!") def disconnect_plugins_database(self, name): if self.local_storage.get("connected_plugins_databases"): if name in self.local_storage.get("connected_plugins_databases"): self.local_storage.delete_element( "connected_plugins_databases", name) self.local_storage.delete_element("plugins", name) return self.badges.output_error("No such plugins database connected!") def connect_payloads_database(self, name, path): if self.local_storage.get("connected_payloads_databases"): if name in self.local_storage.get("connected_payloads_databases"): self.bagdes.output_error( "Payloads database already connected!") return if not os.path.exists(path) or not str.endswith(path, "json"): self.badges.output_error("Not a payloads database!") return try: database = json.load(open(path)) except Exception: self.badges.output_error("Failed to connect payloads database!") return if '__database__' not in database.keys(): self.badges.output_error("No __database__ section found!") return if database['__database__']['type'] != "payloads": self.badges.output_error("Not a payloads database!") return del database['__database__'] payloads = {name: database} data = {name: {'path': path}} if not self.local_storage.get("connected_payloads_databases"): self.local_storage.set("connected_payloads_databases", dict()) self.local_storage.update("connected_payloads_databases", data) if self.local_storage.get("payloads"): self.local_storage.update("payloads", payloads) else: self.local_storage.set("payloads", payloads) def connect_modules_database(self, name, path): if self.local_storage.get("connected_modules_databases"): if name in self.local_storage.get("connected_modules_databases"): self.bagdes.output_error("Modules database already connected!") return if not os.path.exists(path) or not str.endswith(path, "json"): self.badges.output_error("Not a modules database!") return try: database = json.load(open(path)) except Exception: self.badges.output_error("Failed to connect modules database!") return if '__database__' not in database.keys(): self.badges.output_error("No __database__ section found!") return if database['__database__']['type'] != "modules": self.badges.output_error("Not a modules database!") return del database['__database__'] modules = {name: database} data = {name: {'path': path}} if not self.local_storage.get("connected_modules_databases"): self.local_storage.set("connected_modules_databases", dict()) self.local_storage.update("connected_modules_databases", data) if self.local_storage.get("modules"): self.local_storage.update("modules", modules) else: self.local_storage.set("modules", modules) def connect_plugins_database(self, name, path): if self.local_storage.get("connected_plugins_databases"): if name in self.local_storage.get("connected_plugins_databases"): self.bagdes.output_error("Plugins database already connected!") return if not os.path.exists(path) or not str.endswith(path, "json"): self.badges.output_error("Not a database!") return try: database = json.load(open(path)) except Exception: self.badges.output_error("Failed to connect plugins database!") return if '__database__' not in database.keys(): self.badges.output_error("No __database__ section found!") return if database['__database__']['type'] != "plugins": self.badges.output_error("Not a plugins database!") return del database['__database__'] plugins = {name: database} data = {name: {'path': path}} if not self.local_storage.get("connected_plugins_databases"): self.local_storage.set("connected_plugins_databases", dict()) self.local_storage.update("connected_plugins_databases", data) if self.local_storage.get("plugins"): self.local_storage.update("plugins", plugins) else: self.local_storage.set("plugins", plugins)
class Jobs: def __init__(self): self.exceptions = Exceptions() self.tables = Tables() self.badges = Badges() self.local_storage = LocalStorage() self.modules = Modules() self.job_process = None def stop_dead(self): jobs = self.local_storage.get("jobs") if jobs: for job_id in list(jobs): if not jobs[job_id]['job_process'].is_alive(): self.delete_job(job_id) def check_jobs(self): if not self.local_storage.get("jobs"): return True return False def check_module_job(self, module_name): jobs = self.local_storage.get("jobs") if jobs: for job_id in jobs.keys(): if jobs[job_id]['module_name'] == module_name: return True return False def exit_jobs(self): if self.check_jobs(): return True self.badges.output_warning("You have some running jobs.") if self.badges.input_question("Exit anyway? [y/N] ").lower() in [ 'yes', 'y' ]: self.badges.output_process("Stopping all jobs...") self.stop_all_jobs() return True return False def stop_all_jobs(self): if not self.check_jobs(): for job_id in list(self.local_storage.get("jobs").keys()): self.delete_job(job_id) def stop_job(self, job): if job.is_alive(): exc = ctypes.py_object(SystemExit) res = ctypes.pythonapi.PyThreadState_SetAsyncExc( ctypes.c_long(job.ident), exc) if res == 0: raise self.exceptions.GlobalException if res > 1: ctypes.pythonapi.PyThreadState_SetAsyncExc(job.ident, None) raise self.exceptions.GlobalException def start_job(self, job_function, job_arguments): self.job_process = threading.Thread(target=job_function, args=job_arguments) self.job_process.setDaemon(True) self.job_process.start() def delete_job(self, job_id): if not self.check_jobs(): job_id = int(job_id) if job_id in list(self.local_storage.get("jobs").keys()): try: self.stop_job( self.local_storage.get("jobs")[job_id]['job_process']) self.local_storage.delete_element("jobs", job_id) except Exception: self.badges.output_error("Failed to stop job!") else: self.badges.output_error("Invalid job id!") else: self.badges.output_error("Invalid job id!") def create_job(self, job_name, module_name, job_function, job_arguments=()): self.start_job(job_function, job_arguments) if not self.local_storage.get("jobs"): self.local_storage.set("jobs", dict()) job_id = len(self.local_storage.get("jobs")) job_data = { job_id: { 'job_name': job_name, 'module_name': module_name, 'job_process': self.job_process } } self.local_storage.update("jobs", job_data) return job_id