def _acquire_lock(self): LOCK_TIMEOUT = timedelta(minutes=120) WAIT_STEP = 15 vms = Internals.get(name='virtual_machines') if vms is None: vms = Internals({'name': 'virtual_machines'}) vms.save() locked_vm = False while not locked_vm: for i, label in enumerate(self.labels): self._use_vm(i) last_locked = "{}.last_locked".format(self.vm_record) if vms.update_value([self.vm_record, 'locked'], True): vms.update_value([self.vm_record, 'last_locked'], datetime.now()) locked_vm = True break expired_date = datetime.now() - LOCK_TIMEOUT if vms._update({'$set': { last_locked: datetime.now() }}, {last_locked: { '$lt': expired_date }}): vms.update_value([self.vm_record, 'locked'], True) locked_vm = True break if not locked_vm: sleep(WAIT_STEP)
def do_pull(self, path=None): print "[+] Pulling '{}'".format(self['name']) try: repo = Repo(path or self.path()) if self['private']: with repo.git.custom_environment( GIT_SSH_COMMAND=self['ssh_cmd']): repo.remotes.origin.pull() else: repo.remotes.origin.pull() # Make sure we delete orphan .pyc files for root, dirs, files in os.walk(path or self.path()): for f in files: f = os.path.join(root, f) if f.endswith(".pyc") and not os.path.exists(f[:-1]): print "Deleting orphan file '{}'".format(f) os.remove(f) updates = Internals.get(name="updates") updates.update_value("last_update", time()) return True except Exception, e: self['status'] = 'error' self['error_msg'] = 'Could not update repository.\n{}'.format(e) self.save() return False
def start(self): try: self.last_run = time() self.clean_temp_dir() self.update_modules() self.process = self._new_celery_worker() while True: updates = Internals.get(name='updates') if updates['last_update'] > self.last_run: # Stop running worker os.kill(self.process.pid, signal.SIGTERM) self.process.wait() # Update modules if needed self.update_modules() # Restart worker self.process = self._new_celery_worker() self.last_run = time() sleep(self.refresh_interval) except KeyboardInterrupt: not_finished = True while not_finished: try: self.process.wait() not_finished = False except KeyboardInterrupt: pass
def _acquire_lock(self): LOCK_TIMEOUT = timedelta(minutes=120) WAIT_STEP = 15 vms = Internals.get(name='virtual_machines') if vms is None: vms = Internals({'name': 'virtual_machines'}) vms.save() last_locked = "{}.last_locked".format(self.vm_record) while True: if vms.update_value([self.vm_record, 'locked'], True): vms.update_value([self.vm_record, 'last_locked'], datetime.now()) break expired_date = datetime.now() - LOCK_TIMEOUT if vms._update({'$set': { last_locked: datetime.now() }}, {last_locked: { '$lt': expired_date }}): vms.update_value([self.vm_record, 'locked'], True) break sleep(WAIT_STEP)
def create_internals(): updates = Internals.get(name='updates') if updates is None: updates = Internals({ 'name': 'updates', 'last_update': time() }) updates.save()
class Repository(MongoDict): collection_name = 'repositories' def __init__(self, values={}): keyfile = os.path.join(FAME_ROOT, "conf", "id_rsa") self['ssh_cmd'] = "ssh -o StrictHostKeyChecking=no -i {}".format( keyfile) MongoDict.__init__(self, values) def delete(self): # First, remove modules from database for module in ModuleInfo.find(): if module['path'].startswith('fame.modules.{}.'.format( self['name'])): module.delete() # Then, delete the files try: rmtree(self.path()) except: pass # Finally, delete record of repository MongoDict.delete(self) def path(self): return os.path.join(FAME_ROOT, 'fame', 'modules', self['name']) def clone(self): clone_repository.apply_async((self['_id'], ), queue='updates') #clone_repository(self['_id']) def do_clone(self): print "[+] Cloning '{}'".format(self['name']) try: if self['private']: Repo.clone_from(self['address'], self.path(), env=dict(GIT_SSH_COMMAND=self['ssh_cmd'])) else: Repo.clone_from(self['address'], self.path()) dispatcher.update_modules(self) self.update_value('status', 'active') except Exception, e: self['status'] = 'error' self[ 'error_msg'] = 'Could not clone repository, probably due to authentication issues.\n{}'.format( e) self.save() internals = Internals.get(name="updates") internals.update_value("last_update", time())
def do_clone(self): print(("[+] Cloning '{}'".format(self['name']))) try: if self['private']: Repo.clone_from(self['address'], self.path(), env=dict(GIT_SSH_COMMAND=self['ssh_cmd'])) else: Repo.clone_from(self['address'], self.path()) dispatcher.update_modules(self) self.update_value('status', 'active') except Exception as e: self['status'] = 'error' self['error_msg'] = 'Could not clone repository, probably due to authentication issues.\n{}'.format(e) self.save() internals = Internals.get(name="updates") internals.update_value("last_update", time())
def _restore_vm(self): if self.always_ready: vms = Internals.get(name='virtual_machines') if self.name not in vms[self.vm_record]: vms.update_value([self.vm_record, self.name], 1) else: vms.update_value([self.vm_record, self.name], vms[self.vm_record][self.name] + 1) if vms[self.vm_record][self.name] >= self.restore_after: self.should_restore = True if self.should_restore: self._vm.restore() vms.update_value([self.vm_record, self.name], 0) else: self._vm.stop()
def do_clone(self, path=None): print "[+] Cloning '{}'".format(self['name']) try: if self['private']: Repo.clone_from(self['address'], path or self.path(), env=dict(GIT_SSH_COMMAND=self['ssh_cmd'])) else: Repo.clone_from(self['address'], path or self.path()) internals = Internals.get(name="updates") internals.update_value("last_update", time()) return True except Exception, e: self['status'] = 'error' self[ 'error_msg'] = 'Could not clone repository, probably due to authentication issues.\n{}'.format( e) # noqa self.save() return False
def _release_lock(self): vms = Internals.get(name='virtual_machines') vms.update_value([self.vm_record, 'locked'], False)
def do_pull(self): print "[+] Pulling '{}'".format(self['name']) try: repo = Repo(self.path()) if self['private']: with repo.git.custom_environment( GIT_SSH_COMMAND=self['ssh_cmd']): repo.remotes.origin.pull() else: repo.remotes.origin.pull() # Make sure we delete orphan .pyc files for root, dirs, files in os.walk(self.path()): for f in files: f = os.path.join(root, f) if f.endswith(".pyc") and not os.path.exists(f[:-1]): print "Deleting orphan file '{}'".format(f) os.remove(f) dispatcher.update_modules(self) self.update_value('status', 'active') except Exception, e: self['status'] = 'error' self['error_msg'] = 'Could not update repository.\n{}'.format(e) self.save() updates = Internals.get(name="updates") updates.update_value("last_update", time())