def store_service_config(name, store_config=True, config_overrides=None, enabled=True): cls = class_by_name(name) if '.' in name else service_by_name(name) if not hasattr(cls, "get_default_config"): raise Exception(name + " is not an AL service. Make sure the class path you've entered is valid.") cfg = cls.get_default_config() cfg['enabled'] = enabled if '.' in name: cfg['classpath'] = name if config_overrides: for cfg_key, cfg_value in config_overrides.iteritems(): if cfg_key not in cfg['config'] and cfg_key != 'PLUMBER_MAX_QUEUE_SIZE': raise Exception("Config override %s is not a valid configuration option for %s" % (cfg_key, name)) cfg['config'][cfg_key] = cfg_value if store_config: srv_config = forge.get_datastore().get_service(cfg['name']) if srv_config: srv_config.update(cfg) else: srv_config = cfg forge.get_datastore().save_service(cfg['name'], srv_config) return srv_config return cfg
def validate_rules(self, rulefile, datastore=False): change = False while True: try: self.paranoid_rule_check(rulefile) return change # If something goes wrong, clean rules until valid file given except Exception as e: change = True if e.message.startswith('yara.SyntaxError'): e_line = int(e.message.split('):', 1)[0].split("(", -1)[1]) e_message = e.message.split("): ", 1)[1] try: invalid_rule, reline = self.clean(rulefile, e_line, e_message) except Exception as ve: raise ve # If datastore object given, change status of signature to INVALID in Riak if datastore: from assemblyline.al.common import forge store = forge.get_datastore() config = forge.get_config() signature_user = config.services.master_list.Yara.config.SIGNATURE_USER # Get the offending sig ID sig_query = "name:{} AND meta.al_status:(DEPLOYED OR NOISY)".format(invalid_rule) sigl = store.list_filtered_signature_keys(sig_query) # Mark and update Riak store = forge.get_datastore() for sig in sigl: sigdata = store.get_signature(sig) # Check this in case someone already marked it as invalid try: if sigdata['meta']['al_status'] == 'INVALID': continue except KeyError: pass sigdata['meta']['al_status'] = 'INVALID' today = datetime.date.today().isoformat() sigdata['meta']['al_state_change_date'] = today sigdata['meta']['al_state_change_user'] = signature_user sigdata['comments'].append("AL ERROR MSG:{0}. Line:{1}".format(e_message.rstrip().strip(), reline)) store.save_signature(sig, sigdata) else: raise e continue
def install(alsi): from copy import deepcopy from assemblyline.al.common import forge from assemblyline.common.user_defaults import ACCOUNT_DEFAULT, SETTINGS_DEFAULT mm_user = alsi.config['core']['middleman']['user'] ds = forge.get_datastore() mm_user_data = ds.get_user(mm_user) if not mm_user_data: mm_user_data = deepcopy(ACCOUNT_DEFAULT) mm_user_data['api_quota'] = 256 mm_user_data['classification'] = alsi.config['core']['middleman'][ 'classification'] mm_user_data['groups'] = ["MIDDLEMAN"] mm_user_data['name'] = mm_user mm_user_data['uname'] = mm_user ds.save_user(mm_user, mm_user_data) mm_options = deepcopy(SETTINGS_DEFAULT) mm_options['classification'] = alsi.config['core']['middleman'][ 'classification'] ds.save_user(mm_user + "_options", mm_options) alsi.sudo_install_file( 'assemblyline/al/install/etc/init/middleman_instance.conf', '/etc/init/middleman_instance.conf') alsi.sudo_install_file('assemblyline/al/install/etc/init/middleman.conf', '/etc/init/middleman.conf') if not os.path.exists('/etc/init.d/middleman'): alsi.runcmd('sudo ln -s /lib/init/upstart-job /etc/init.d/middleman') alsi.milestone("middleman install complete")
def add_to_profile(profile_name, service_key): ds = forge.get_datastore() profile = ds.get_profile(profile_name) if not profile: raise Exception("Could not find profile: %s") profile['services'][service_key] = {'service_overrides': {}, 'workers': 1} ds.save_profile(profile_name, profile)
def main(bucket_list, journal_queues): ds = forge.get_datastore() queues = { x: queue.NamedQueue('d-%s' % x, db=DATABASE_NUM) for x in set(journal_queues).union(set(bucket_list)) } Thread(target=track_status, name="queues_status", args=(queues, )).start() log.info("Ready!") loader_threads = { x: Thread(target=load_expired, name="loader_%s" % x, args=(ds, x, queues[x])) for x in bucket_list } loader_threads.update({ 'journal_%s' % x: Thread(target=load_journal, name="journal_loader_%s" % x, args=(x, queues[x])) for x in journal_queues }) for thread in loader_threads.itervalues(): thread.start() for thread in loader_threads.itervalues(): thread.join()
def start(self): self.vmm = CuckooVmManager(self.cfg) self.cm = CuckooContainerManager(self.cfg, self.vmm) self._register_cleanup_op({ 'type': 'shell', 'args': shlex.split("docker rm --force %s" % self.cm.name) }) self.log.debug("VMM and CM started!") # Start the container self.cuckoo_ip = self.cm.start_container(self.cm.name) self.restart_interval = random.randint(45, 55) self.file_name = None self.set_urls() self.ssdeep_match_pct = int(self.cfg.get("dedup_similar_percent", 80)) for param in forge.get_datastore().get_service( self.SERVICE_NAME)['submission_params']: if param['name'] == "routing": self.enabled_routes = param['list'] if self.enabled_routes[0] != param['default']: self.enabled_routes.remove(param['default']) self.enabled_routes.insert(0, param['default']) if self.enabled_routes is None: raise ValueError("No routing submission_parameter.") self.log.debug("Cuckoo started!")
def ingester(): # df node def # pylint:disable=R0912 datastore = forge.get_datastore() user_groups = {} # Move from ingest to unique and waiting queues. # While there are entries in the ingest queue we consume chunk_size # entries at a time and move unique entries to uniqueq / queued and # duplicates to their own queues / waiting. while running: while True: result = completeq.pop(blocking=False) # df pull pop if not result: break completed(Task(result)) # df push calls entry = ingestq.pop(timeout=1) # df pull pop if not entry: continue trafficq.push(entry) # df push push sha256 = entry.get('sha256', '') if not sha256 or len(sha256) != 64: logger.error("Invalid sha256: %s", entry) continue entry['md5'] = entry.get('md5', '').lower() entry['sha1'] = entry.get('sha1', '').lower() entry['sha256'] = sha256.lower() ingest(datastore, user_groups, entry) # df push calls datastore.close()
def undrain(self, msg): self.store = forge.get_datastore() if self.service_manager: self.service_manager.undrain() if self.vm_manager: self.vm_manager.undrain() return True
def __init__(self, wid, worker_type, working_dir, instance_id): self.working_dir = working_dir self.worker_id = wid self.ds = forge.get_datastore() self.worker_type = worker_type self.instance_id = instance_id if worker_type == TYPE_BACKUP: self.hash_queue = remote_datatypes.Hash("r-hash_%s" % self.instance_id, db=DATABASE_NUM) self.follow_queue = queue.NamedQueue("r-follow_%s" % self.instance_id, db=DATABASE_NUM, ttl=1800) self.queue = queue.NamedQueue("r-backup_%s" % self.instance_id, db=DATABASE_NUM, ttl=1800) self.done_queue = queue.NamedQueue("r-backup-done_%s" % self.instance_id, db=DATABASE_NUM, ttl=1800) else: self.hash_queue = None self.follow_queue = None self.queue = None self.done_queue = queue.NamedQueue("r-restore-done_%s" % self.instance_id, db=DATABASE_NUM, ttl=1800)
def install(alsi=None): from assemblyline.al.common import forge ds = forge.get_datastore() ip = net.get_hostip() mac = net.get_mac_for_ip(ip) existing_reg = ds.get_node(mac) if existing_reg: alsi.info("Registration already exist. Skipping...") return reg = DEFAULT_CORE_REGISTRATION.copy() reg['hostname'] = net.get_hostname() reg['ip'] = ip reg['mac_address'] = mac reg['machine_info'] = sysinfo.get_machine_info() reg['platform'] = sysinfo.get_platform() if 'roles' not in reg: reg['roles'] = [] if "dispatcher" not in reg["roles"]: reg['roles'].append("dispatcher") if "middleman" not in reg["roles"]: reg['roles'].append("middleman") ds.save_node(mac, reg) alsi.info("Core server registered!")
def main(): log.init_logging('run_service') if len(sys.argv) != 2: usage() exit(1) name = sys.argv[1] try: svc_class = class_by_name(name) if '.' in name else service_by_name( name) except: print 'Could not load service "%s".\n%s' % (name, get_valid_service_list()) raise logger.info('Running service in stand-alone mode. CTRL-C to exit.') # noinspection PyBroadException try: cfg = forge.get_datastore().get_service(svc_class.SERVICE_NAME).get( "config", {}) except: # pylint: disable=W0702 cfg = {} service_driver = ServiceDriver(svc_class, cfg, 86400, NUM_WORKERS) service_driver.start() try: while True: send_minimal_heartbeat(svc_class.SERVICE_NAME, NUM_WORKERS) time.sleep(config.system.update_interval) except KeyboardInterrupt: print 'Exiting.' finally: service_driver.stop_hard()
def main(shard): log.init_logging('dispatcher') ds = forge.get_datastore() service_proxies = ServiceProxyManager(ds.list_service_keys()) dispatcher = Dispatcher(service_proxies, shard=shard, debug=False) dispatcher.start()
def __init__(self, vmname, cfg=None): self.ds = forge.get_datastore() self.vm_name = vmname + '.001' if cfg: self.vm_cfg = cfg.workers.virtualmachines.master_list.get(vmname, {}).get('cfg', None) else: self.vm_cfg = self.ds.get_virtualmachine(vmname) if not self.vm_cfg: raise Exception("Could not find VM %s in the seed" % vmname) self.vmm = libvirt.open(None)
def __init__(self, tasker_id_p): self.action_queue = queue.PriorityQueue('alert-actions', db=DATABASE_NUM) self.worker_queues_map = { x: queue.PriorityQueue('alert-actions-worker-%s' % x, db=DATABASE_NUM) for x in range(WORKER_COUNT) } self.datastore = forge.get_datastore() self.tasker_id = tasker_id_p
def __init__(self, logger=None): if not logger: from assemblyline.al.common import log as al_log al_log.init_logging('yara_importer') logger = logging.getLogger('assemblyline.yara_importer') logger.setLevel(logging.INFO) yara_parser_class = forge.get_yara_parser() self.ds = forge.get_datastore() self.yp = yara_parser_class() self.log = logger self._id_cache = {} self._name_cache = []
def __init__(self, server_url=None, datastore=None): if not server_url: server_url = config.submissions.url self.server_url = server_url self.transport = forge.get_filestore() self.datastore = datastore self.is_unix = os.name == "posix" if not self.is_unix: from assemblyline_client import Client self.client = Client(self.server_url, auth=SUBMISSION_AUTH) elif self.datastore is None: self.datastore = forge.get_datastore()
def __init__(self, backup_file_path): self.backup_file_path = backup_file_path self.ds = forge.get_datastore() # Static maps self.BUCKET_MAP = { "blob": self.ds.blobs, "node": self.ds.nodes, "profile": self.ds.profiles, "signature": self.ds.signatures, "user": self.ds.users, } self.VALID_BUCKETS = sorted(self.BUCKET_MAP.keys())
def main(): if len(sys.argv) == 1: print "Usage: %s <One or more prepared VM tarballs>" sys.exit(7) try: svc_class = service_by_name("Cuckoo") except: print 'Could not load service "%s".\n' \ 'Valid options:\n%s' % ("Cuckoo", [s['name'] for s in forge.get_datastore().list_services()]) sys.exit(7) cfg = forge.get_datastore().get_service(svc_class.SERVICE_NAME).get( "config", {}) config = forge.get_config() local_meta_root = os.path.join(config.system.root, cfg['REMOTE_DISK_ROOT']) vm_meta_path = os.path.join(local_meta_root, cfg['vm_meta']) out_config = vm_meta_path out_directory = os.path.dirname(out_config) vm_list = sys.argv[1:] cuckoo_config = [] for vm in vm_list: for js in install_vm_meta(out_directory, vm, ['']): cuckoo_config.append(js) with open(out_config, "w") as fh: json.dump(cuckoo_config, fh, sort_keys=True, indent=4, separators=(',', ': ')) print "Wrote %i Definitions to %s!" % (len(cuckoo_config), out_config)
def start(self): self._state = State.STARTING self.log.info("VmManager starting on (%s:%s).", self.host_ip, self.host_mac) with self.vmm_lock: self.vmm = libvirt.open(None) self.vmrevert_scheduler = apscheduler.scheduler.Scheduler() self.store = forge.get_datastore() # clean any state that might be left from previous run. self._destroy_all() # install and start vms self._reconcile_config() self.log.info("Starting all instances.") self._start_all() self._schedule_automatic_reverts() self._state = State.RUNNING
def __init__(self, vmcfg=None): self.disk_root = config.workers.virtualmachines.disk_root if not os.path.exists(self.disk_root): os.makedirs(self.disk_root) self.vmm = None self.cfg = vmcfg if vmcfg is None: self.cfg = get_vmcfg_for_localhost() self.vmrevert_scheduler = None self.host_ip = net.get_hostip() self.host_mac = net.get_mac_for_ip(self.host_ip) self.log = logging.getLogger('assemblyline.vmm') self.vm_profiles = {} self.vmm_lock = threading.Lock() self._state = State.INIT self._hostagent_client = ServiceAgentClient(async=True) self.store = forge.get_datastore() self.vm_configs = {}
def reregister_services(store_config=True, config_overrides=None): failed = [] passed = [] services = forge.get_datastore().list_services() for svc in services: class_name = svc.get('classpath', "al_services.%s.%s" % (svc['repo'], svc['class_name'])) pprint(class_name) try: store_service_config(class_name, store_config, config_overrides) passed.append(class_name) except ImportError as ie: failed.append((class_name, ie)) logging.error('Existing service cannot be imported for inspection: %s. Skipping.', class_name) logging.info("Succeeded: \n\t%s", "\n\t".join(passed)) logging.info("Failed:") for (c, e) in failed: logging.info("\t%s - %s", c, e)
def __init__(self): self.mac = net.get_mac_address() self.store = forge.get_datastore() self.log = logging.getLogger('assemblyline.control') self.log.info('Starting Controller: MAC[%s] STORE[%s]' % (self.mac, self.store)) # This hosts registration from riak (Hosts tab in UI). self.jobs = LocalQueue() self.last_heartbeat = 0 self.rpc_handlers = { ControllerRequest.HEARTBEAT: self.controller_heartbeat, ControllerRequest.START: self.hostagent_start, ControllerRequest.STOP: self.hostagent_stop, ControllerRequest.RESTART: self.hostagent_restart, ControllerRequest.STATUS: self.hostagent_status, } self._should_run = True self.executor_thread = None self.heartbeat_thread = None
def main(): log.init_logging('test') logger = logging.getLogger('assemblyline.test') store = forge.get_datastore() sids = [] for x in store.stream_search( 'submission', 'times.completed:[2015-01-30T00:00:00.000Z TO 2015-01-30T00:59:59.999Z]' ): sid = x['submission.sid'] sids.append(sid) count = 0 submissions = store.get_submissions(sids) for submission in submissions: if submission.get('state', '') != 'completed': continue if len(submission['files']) != 1: continue _, srl = submission['files'][0] fileinfo = store.get_file(srl) if not fileinfo: continue submission = submission.get('submission', {}) if not submission: continue metadata = submission.get('metadata', {}) if not metadata: continue metadata['ignore_submission_cache'] = True metadata['ignore_cache'] = False metadata['md5'] = fileinfo['md5'] metadata['sha1'] = fileinfo['sha1'] metadata['sha256'] = fileinfo['sha256'] metadata['size'] = fileinfo['size'] resubmit(metadata) count += 1 if count >= 1000: break logger.info('Resubmitted %d submissions for testing', count)
def __init__(self): self.ip = net.get_hostip() self.mac = net.get_mac_for_ip(self.ip) self.store = forge.get_datastore() self.log = logging.getLogger('assemblyline.agent') self.log.info('Starting HostAgent: MAC[%s] STORE[%s]' % (self.mac, self.store)) # This hosts registration from riak (Hosts tab in UI). self.registration = None self.service_manager = None self.vm_manager = None self.flex_manager = None self.lock = None self.consumer_thread = None self._should_run = False self.host_profile = {} self.executor_thread = None # Chores are actions that we run periodically and which we coallesce # when the same chore is requested multiple times in the same tick. # Jobs are executed as they are received. self.jobs = LocalQueue() self.last_heartbeat = 0 self.rpc_handlers = { AgentRequest.PING: self.ping, AgentRequest.DRAIN: self.drain, AgentRequest.UNDRAIN: self.undrain, AgentRequest.SHUTDOWN: self.shutdown, AgentRequest.VM_LIST: self.list_vms, AgentRequest.VM_START: self.start_vm, AgentRequest.VM_STOP: self.stop_vm, AgentRequest.VM_STOP_ALL: self.stop_all_vms, AgentRequest.VM_RESTART: self.restart_vm, AgentRequest.VM_REFRESH_ALL: self.refresh_vm_all, AgentRequest.VM_REFRESH_FLEET: self.refresh_vm_fleet, AgentRequest.VM_GET_REVERT_TIMES: self.vm_get_revert_times, AgentRequest.START_SERVICES: self.start_services, AgentRequest.STOP_SERVICES: self.stop_services, } self._should_run = True
def __init__(self, working_dir, worker_count=50, spawn_workers=True): self.working_dir = working_dir self.ds = forge.get_datastore() self.plist = [] self.instance_id = str(uuid.uuid4()) self.follow_queue = queue.NamedQueue("r-follow_%s" % self.instance_id, db=DATABASE_NUM, ttl=1800) self.hash_queue = remote_datatypes.Hash("r-hash_%s" % self.instance_id, db=DATABASE_NUM) self.backup_queue = queue.NamedQueue('r-backup_%s' % self.instance_id, db=DATABASE_NUM, ttl=1800) self.backup_done_queue = queue.NamedQueue("r-backup-done_%s" % self.instance_id, db=DATABASE_NUM, ttl=1800) self.restore_done_queue = queue.NamedQueue("r-restore-done_%s" % self.instance_id, db=DATABASE_NUM, ttl=1800) self.bucket_error = [] self.BUCKET_MAP = { "alert": self.ds.alerts, "blob": self.ds.blobs, "emptyresult": self.ds.emptyresults, "error": self.ds.errors, "file": self.ds.files, "filescore": self.ds.filescores, "node": self.ds.nodes, "profile": self.ds.profiles, "result": self.ds.results, "signature": self.ds.signatures, "submission": self.ds.submissions, "user": self.ds.users, } self.VALID_BUCKETS = sorted(self.BUCKET_MAP.keys()) self.worker_count = worker_count self.spawn_workers = spawn_workers self.current_type = None
def get_vmcfg_for_localhost(): ip = net.get_hostip() mac = net.get_mac_for_ip(ip) store = forge.get_datastore() host_registration = store.get_node(mac) if not host_registration: raise ConfigException('Could not find host registration fr %s' % mac) profile_name = host_registration.get('profile', None) if not profile_name: raise ConfigException('Could not find profile for host: %s' % mac) host_profile = store.get_profile(profile_name) if not host_profile: raise ConfigException('Could not fetch host profile %s' % profile_name) vm_config = host_profile.get('virtual_machines', None) if not vm_config: raise ConfigException('Could not find virtual machine section in %s' % profile_name) store.client.close() return vm_config
def dropper(): # df node def datastore = forge.get_datastore() while running: raw = dropq.pop(timeout=1) # df pull pop if not raw: continue notice = Notice(raw) send_notification(notice) c12n = notice.get('classification', config.core.middleman.classification) expiry = now_as_iso(86400) sha256 = notice.get('sha256') datastore.save_or_freshen_file(sha256, {'sha256': sha256}, expiry, c12n) datastore.close()
def __init__(self): # Delay these imports so most nodes don't import them. global Scheduler from apscheduler.scheduler import Scheduler self.bottleneck_queue_sizes = {} self.cores = None self.datastore = forge.get_datastore() self.flex_profile = None self.flex_scheduler = None self.log = logging.getLogger('assemblyline.flex') self.mac = net.get_mac_for_ip(net.get_hostip()) self.main_bottleneck = '' self.needs_cleanup = True self.previous_queue_sizes = {} self.safe_start_dict = {} self.safeq = NamedQueue('safe-start-%s' % self.mac) self.service_manager = None self.ram_mb = None self.tick_count = 0 self.vm_manager = None
def get_service_queue_lengths(): global ds # pylint: disable=W0603 if not ds: ds = forge.get_datastore() # Default is to return all services in a dict of class_name: queue_size. queue_lengths = {} services = ds.list_services() for svc in services: # noinspection PyBroadException try: if not svc: continue classpath = svc.get( 'classpath', "al_services.%s.%s" % (svc['repo'], svc['class_name'])) queue_lengths[svc['name']] = get_service_queue_length(classpath) except Exception: # pylint: disable=W0703 log.exception('while getting queue length for %s', svc['name']) return queue_lengths
def main(shard): log.init_logging('dispatcher') logger = logging.getLogger('assemblyline.dispatch') r = redis.StrictRedis(config.core.redis.nonpersistent.host, config.core.redis.nonpersistent.port, config.core.redis.nonpersistent.db) r.delete('ingest-queue-' + shard) store = forge.get_datastore() store.commit_index('submission') query = 'state:submitted AND times.submitted:[NOW-1DAY TO *]' sids = [] for x in store.stream_search('submission', query): sid = x['submission.sid'] if str(forge.determine_dispatcher(sid)) == shard: sids.append(sid) count = 0 submissions = store.get_submissions(sids) for submission in submissions: if submission.get('state', '') != 'submitted': sid = submission.get('sid', '') if sid: store.save_submission(sid, submission) continue submission['request'] = {} for path, srl in submission['files']: submission['fileinfo'] = store.get_file(srl) submission['request']['path'] = path submission['request']['srl'] = srl resubmit(submission) count += 1 logger.info('Resubmitted %d submissions to dispatcher %s.', count, shard)