def _run_work(self): while not self.isfini: self._pool_avail += 1 work = self.workq.get() self._pool_avail -= 1 if work is None: return self.fire('pool:work:init', work=work) task, info = work jid = info.get('jid') try: func, args, kwargs = task ret = func(*args, **kwargs) # optionally generate a job event if jid is not None: self.fire('job:done', jid=jid, ret=ret) except Exception as e: if jid is not None: self.fire('job:done', jid=jid, **s_common.excinfo(e)) self.fire('pool:work:fini', work=work)
def _onTeleOnMesg(self, sock, mesg): try: # set the socket tx method as the callback jid = mesg[1].get('jid') ons = mesg[1].get('ons') name = mesg[1].get('name') item = self.shared.get(name) if item is None: raise s_common.NoSuchObj(name=name) user = sock.get('syn:user') func = getattr(item, 'on', None) if func is None: return sock.tx(s_common.tufo('job:done', jid=jid, ret=False)) self._reqUserAllowed(user, 'tele:call', name, 'on') for evnt, ontups in ons: # (<evnt>, ( (<iden>,<filt>), ... )) onhelp, new = self._getOnHelp(name, evnt) if new: func(evnt, onhelp.dist) for iden, filt in ontups: onhelp.addOnInst(sock, iden, filt) return sock.tx(s_common.tufo('job:done', jid=jid, ret=True)) except Exception as e: sock.tx(s_common.tufo('job:done', jid=jid, **s_common.excinfo(e)))
def _onTeleOffMesg(self, sock, mesg): # set the socket tx method as the callback try: jid = mesg[1].get('jid') evnt = mesg[1].get('evnt') name = mesg[1].get('name') iden = mesg[1].get('iden') item = self.shared.get(name) if item is None: raise s_common.NoSuchObj(name=name) onhelp, new = self._getOnHelp(name, evnt) onhelp.delOnInst(sock, iden) return sock.tx(s_common.tufo('job:done', jid=jid, ret=True)) except Exception as e: errinfo = s_common.excinfo(e) sock.tx( s_common.tufo('job:done', jid=jid, err=errinfo.get('err'), errinfo=errinfo))
def _runTeleCall(self, mesg): jid = mesg[1].get('jid') name = mesg[1].get('name') task = mesg[1].get('task') suid = mesg[1].get('suid') retinfo = dict(suid=suid, jid=jid) try: item = self._tele_pushed.get(name) if item is None: return self._txTeleSock('tele:retn', err='NoSuchObj', errmsg=name, **retinfo) meth, args, kwargs = task func = getattr(item, meth, None) if func is None: return self._txTeleSock('tele:retn', err='NoSuchMeth', errmsg=meth, **retinfo) self._txTeleSock('tele:retn', ret=func(*args, **kwargs), **retinfo) except Exception as e: retinfo.update(s_common.excinfo(e)) return self._txTeleSock('tele:retn', **retinfo)
async def send(self, host, port=25, user=None, passwd=None, usetls=False, starttls=False, timeout=60): self.runt.confirm(('storm', 'inet', 'smtp', 'send')) try: if self.bodytext is None and self.bodyhtml is None: mesg = 'The storm:smtp:message has no HTML or text body.' raise s_exc.StormRuntimeError(mesg=mesg) host = await s_stormtypes.tostr(host) port = await s_stormtypes.toint(port) usetls = await s_stormtypes.tobool(usetls) starttls = await s_stormtypes.tobool(starttls) timeout = await s_stormtypes.toint(timeout) user = await s_stormtypes.tostr(user, noneok=True) passwd = await s_stormtypes.tostr(passwd, noneok=True) message = MIMEMultipart('alternative') if self.bodytext is not None: message.attach(MIMEText(self.bodytext, 'plain', 'utf-8')) if self.bodyhtml is not None: message.attach(MIMEText(self.bodyhtml, 'html', 'utf-8')) for name, valu in self.headers.items(): message[await s_stormtypes.tostr( name)] = await s_stormtypes.tostr(valu) recipients = [await s_stormtypes.tostr(e) for e in self.recipients] futu = aiosmtplib.send(message, port=port, hostname=host, sender=self.sender, recipients=recipients, use_tls=usetls, start_tls=starttls, username=user, password=passwd) await asyncio.wait_for(futu, timeout=timeout) except asyncio.CancelledError: # pragma: no cover raise except Exception as e: return (False, s_common.excinfo(e)) return (True, {})
def _runOperFuncs(self, query, opers): ''' Run the given oper funcs within the query. ''' try: [self._runOperFunc(query, oper) for oper in opers] except Exception as e: query.clear() query.log(excinfo=s_common.excinfo(e))
def __exit__(self, exc, cls, tb): info = { 'sub': self.query.subed, 'add': self.query.added, 'took': s_common.now() - self.stime } if exc is not None: info['excinfo'] = s_common.excinfo(exc) self.query.clear() self.query.log(**info)
def exc(self, exc, **info): ''' Implements the exception log convention for EventBus. A caller is expected to be within the except frame. Args: exc (Exception): The exception to log Returns: None ''' info.update(s_common.excinfo(exc)) self.log(logging.ERROR, str(exc), **info)
def _task_run(self): func, args, kwargs = self._call try: valu = func(*args, **kwargs) if isinstance(valu, types.GeneratorType): for v in valu: self.retn(v) else: self.retn(valu) except Exception as e: self.err(s_common.excinfo(e))
def _runJob(self, job): ''' Actually execute the given job with the caller thread. ''' task = job[1].get('task') if task is None: # TODO This attribute is not set, a bad tufo # sent to _runJob will have unexpected behavior. self.setJobErr(job[0], 'NoJobTask') return try: func, args, kwargs = task ret = func(*args, **kwargs) self.fire('job:done', jid=job[0], ret=ret) except Exception as e: self.fire('job:done', jid=job[0], **s_common.excinfo(e))
def subtask(job): jid = job[0] slot = job[1].get('slot') meld = job[1].get('meld') if meld is not None: s_mindmeld.loadMindMeld(meld) hive = slot[1].get('hive') queen = s_telepath.openurl(job[1].get('queen')) s_scope.set('syn.queen', queen) try: dyntask = job[1].get('dyntask') ret = s_dyndeps.runDynTask(dyntask) queen.tell(hive, 'job:done', jid=jid, ret=ret) except Exception as e: queen.tell(hive, 'job:done', jid=jid, **s_common.excinfo(e))
async def wget(self, url, params=None, headers=None, json=None, body=None, method='GET', ssl=True, timeout=None): ''' Stream a file download directly into the Axon. Args: url (str): The URL to retrieve. params (dict): Additional parameters to add to the URL. headers (dict): Additional HTTP headers to add in the request. json: A JSON body which is included with the request. body: The body to be included in the request. method (str): The HTTP method to use. ssl (bool): Perform SSL verification. timeout (int): The timeout of the request, in seconds. Notes: The response body will be stored, regardless of the response code. The ``ok`` value in the reponse does not reflect that a status code, such as a 404, was encountered when retrieving the URL. The dictionary returned by this may contain the following values:: { 'ok': <boolean> - False if there were exceptions retrieving the URL. 'url': <str> - The URL retrieved (which could have been redirected) 'code': <int> - The response code. 'mesg': <str> - An error message if there was an exception when retrieving the URL. 'headers': <dict> - The response headers as a dictionary. 'size': <int> - The size in bytes of the response body. 'hashes': { 'md5': <str> - The MD5 hash of the response body. 'sha1': <str> - The SHA1 hash of the response body. 'sha256': <str> - The SHA256 hash of the response body. 'sha512': <str> - The SHA512 hash of the response body. } } Returns: dict: An information dictionary containing the results of the request. ''' logger.debug(f'Wget called for [{url}].', extra=await self.getLogExtra(url=s_urlhelp.sanitizeUrl(url))) proxyurl = self.conf.get('http:proxy') cadir = self.conf.get('tls:ca:dir') connector = None if proxyurl is not None: connector = aiohttp_socks.ProxyConnector.from_url(proxyurl) atimeout = aiohttp.ClientTimeout(total=timeout) if ssl is False: pass elif cadir: ssl = s_common.getSslCtx(cadir) else: # default aiohttp behavior ssl = None async with aiohttp.ClientSession(connector=connector, timeout=atimeout) as sess: try: async with sess.request(method, url, headers=headers, params=params, json=json, data=body, ssl=ssl) as resp: info = { 'ok': True, 'url': str(resp.url), 'code': resp.status, 'headers': dict(resp.headers), } hashset = s_hashset.HashSet() async with await self.upload() as upload: async for byts in resp.content.iter_chunked( CHUNK_SIZE): await upload.write(byts) hashset.update(byts) size, _ = await upload.save() info['size'] = size info['hashes'] = dict([(n, s_common.ehex(h)) for (n, h) in hashset.digests()]) return info except asyncio.CancelledError: raise except Exception as e: logger.exception( f'Failed to wget {s_urlhelp.sanitizeUrl(url)}') exc = s_common.excinfo(e) mesg = exc.get('errmsg') if not mesg: mesg = exc.get('err') return { 'ok': False, 'mesg': mesg, }
async def wput(self, sha256, url, params=None, headers=None, method='PUT', ssl=True, timeout=None, filename=None, filemime=None): ''' Stream a blob from the axon as the body of an HTTP request. ''' proxyurl = self.conf.get('http:proxy') cadir = self.conf.get('tls:ca:dir') connector = None if proxyurl is not None: connector = aiohttp_socks.ProxyConnector.from_url(proxyurl) if ssl is False: pass elif cadir: ssl = s_common.getSslCtx(cadir) else: # default aiohttp behavior ssl = None atimeout = aiohttp.ClientTimeout(total=timeout) async with aiohttp.ClientSession(connector=connector, timeout=atimeout) as sess: try: async with sess.request(method, url, headers=headers, params=params, data=self.get(sha256), ssl=ssl) as resp: info = { 'ok': True, 'url': str(resp.url), 'code': resp.status, 'headers': dict(resp.headers), } return info except asyncio.CancelledError: # pramga: no cover raise except Exception as e: logger.exception( f'Error streaming [{sha256}] to [{s_urlhelp.sanitizeUrl(url)}]' ) exc = s_common.excinfo(e) mesg = exc.get('errmsg') if not mesg: mesg = exc.get('err') return { 'ok': False, 'mesg': mesg, }
async def postfiles(self, fields, url, params=None, headers=None, method='POST', ssl=True, timeout=None): ''' Send files from the axon as fields in a multipart/form-data HTTP request. Args: fields (list): List of dicts containing the fields to add to the request as form-data. url (str): The URL to retrieve. params (dict): Additional parameters to add to the URL. headers (dict): Additional HTTP headers to add in the request. method (str): The HTTP method to use. ssl (bool): Perform SSL verification. timeout (int): The timeout of the request, in seconds. Notes: The dictionaries in the fields list may contain the following values:: { 'name': <str> - Name of the field. 'sha256': <str> - SHA256 hash of the file to submit for this field. 'value': <str> - Value for the field. Ignored if a sha256 has been specified. 'filename': <str> - Optional filename for the field. 'content_type': <str> - Optional content type for the field. 'content_transfer_encoding': <str> - Optional content-transfer-encoding header for the field. } The dictionary returned by this may contain the following values:: { 'ok': <boolean> - False if there were exceptions retrieving the URL. 'err': <str> - An error message if there was an exception when retrieving the URL. 'url': <str> - The URL retrieved (which could have been redirected) 'code': <int> - The response code. 'body': <bytes> - The response body. 'headers': <dict> - The response headers as a dictionary. } Returns: dict: An information dictionary containing the results of the request. ''' proxyurl = self.conf.get('http:proxy') cadir = self.conf.get('tls:ca:dir') connector = None if proxyurl is not None: connector = aiohttp_socks.ProxyConnector.from_url(proxyurl) if ssl is False: pass elif cadir: ssl = s_common.getSslCtx(cadir) else: # default aiohttp behavior ssl = None atimeout = aiohttp.ClientTimeout(total=timeout) async with aiohttp.ClientSession(connector=connector, timeout=atimeout) as sess: try: data = aiohttp.FormData() data._is_multipart = True for field in fields: sha256 = field.get('sha256') if sha256: valu = self.get(s_common.uhex(sha256)) else: valu = field.get('value') data.add_field(field.get('name'), valu, content_type=field.get('content_type'), filename=field.get('filename'), content_transfer_encoding=field.get( 'content_transfer_encoding')) async with sess.request(method, url, headers=headers, params=params, data=data, ssl=ssl) as resp: info = { 'ok': True, 'url': str(resp.url), 'code': resp.status, 'body': await resp.read(), 'headers': dict(resp.headers), } return info except asyncio.CancelledError: # pramga: no cover raise except Exception as e: logger.exception( f'Error POSTing files to [{s_urlhelp.sanitizeUrl(url)}]') exc = s_common.excinfo(e) mesg = exc.get('errmsg') if not mesg: mesg = exc.get('err') return { 'ok': False, 'err': mesg, 'url': url, 'body': b'', 'code': -1, 'headers': dict(), }
def _onTeleCallMesg(self, sock, mesg): # tele:call - call a method on a shared object jid = mesg[1].get('jid') sid = mesg[1].get('sid') # check if the socket knows about their auth # ( most likely via SSL client cert ) user = sock.get('syn:user') with s_scope.enter({ 'dmon': self, 'sock': sock, 'syn:user': user, 'syn:auth': self.auth }): try: name = mesg[1].get('name') item = self.shared.get(name) if item is None: # is it a pushed object? pushsock = self.pushed.get(name) if pushsock is not None: # pass along how to reply mesg[1]['suid'] = sock.iden return pushsock.tx(mesg) raise s_common.NoSuchObj(name) task = mesg[1].get('task') meth, args, kwargs = task self._reqUserAllowed(user, 'tele:call', name, meth) func = getattr(item, meth, None) if func is None: raise s_common.NoSuchMeth(meth) if getattr(func, '_tele_clientside', False): name = s_reflect.getMethName(func) raise s_common.TeleClientSide(name=name) ret = func(*args, **kwargs) # handle generator returns specially if isinstance(ret, types.GeneratorType): iden = s_common.guid() txwait = threading.Event() # start off set... txwait.set() self._dmon_yields.add(iden) sock.tx( s_common.tufo('tele:yield:init', jid=jid, iden=iden)) # FIXME opt maxsize = 100000000 def ontxsize(m): size = m[1].get('size') if size >= maxsize: txwait.clear() else: txwait.set() try: sock.onfini(txwait.set) sock.on('sock:tx:size', ontxsize) for item in ret: txwait.wait() # check if we woke due to fini if sock.isfini: break sock.tx( s_common.tufo('tele:yield:item', iden=iden, item=item)) if iden not in self._dmon_yields: break finally: sock.off('sock:tx:size', ontxsize) self._dmon_yields.discard(iden) sock.tx(s_common.tufo('tele:yield:fini', iden=iden)) return sock.tx(s_common.tufo('job:done', jid=jid, ret=ret)) except Exception as e: sock.tx( s_common.tufo('job:done', jid=jid, **s_common.excinfo(e)))
async def _destIterLyrNodeedits(self, writer, queue, lyriden): ''' Batch available source splices in a queue as nodeedits and push to the destination layer proxy. Nodeedit boundaries are defined by the ndef and prov iden. Will run as long as queue is not fini'd. Args: writer (function): Async function that takes (nodeedits, meta) as input queue (s_queue.Window): Layer queue for splices lyriden (str): Layer iden ''' fair_iter = self.push_fair_iter offs_logging = self.offs_logging err_lim = self.err_lim evnt = self._push_evnts[lyriden] ndef = None prov = None nodesplices = [] nodespliceoffs = [] cnt = 0 errs = 0 async for offs, splice in queue: evnt.clear() queuelen = len(queue.linklist) next_ndef = splice[1]['ndef'] next_prov = splice[1].get('prov') # current splice is a new node or has new prov iden or the queue is empty # so create prior node nodeedit and push to destination layer if ndef is not None and (next_ndef != ndef or (prov is not None and next_prov != prov) or queuelen == 0): err, ne, meta = None, None, None try: err, ne, meta = await self._trnNodeSplicesToNodeedit( ndef, nodesplices) if err is None: await writer([ne], meta) self.push_offs.set(lyriden, offs + 1) except asyncio.CancelledError: # pragma: no cover raise except (ConnectionError, s_exc.IsFini): # put back last and nodesplices queue.linklist.appendleft((offs, splice)) qadd = zip(reversed(nodespliceoffs), reversed(nodesplices)) queue.linklist.extendleft(qadd) raise except Exception as e: err = { 'mesg': s_common.excinfo(e), 'splices': nodesplices, 'nodeedits': ne, 'meta': meta } logger.exception(err['mesg']) if err is not None: errs += 1 await self._setLyrErr(lyriden, offs, err) if err_lim != -1 and errs >= err_lim: logger.error(f'Error limit reached') raise SyncErrLimReached( mesg= 'Error limit reached - correct or increase err_lim to continue' ) nodesplices = [] nodespliceoffs = [] ndef = next_ndef prov = next_prov nodesplices.append(splice) nodespliceoffs.append(offs) cnt += 1 if queuelen != 0 and queuelen % 10000 == 0: # pragma: no cover logger.debug( f'{lyriden} queue reader status: read={cnt}, errs={errs}, size={queuelen}' ) if offs % offs_logging == 0 and offs != 0: # pragma: no cover logger.info( f'Destination layer splice push at offset {offs} with {errs} errors' ) if queuelen == 0: evnt.set() await asyncio.sleep(0) elif cnt % fair_iter == 0: await asyncio.sleep(0)
async def wget(self, url, params=None, headers=None, json=None, body=None, method='GET', ssl=True, timeout=None): ''' Stream a file download directly into the axon. ''' connector = None proxyurl = self.conf.get('http:proxy') if proxyurl is not None: connector = aiohttp_socks.ProxyConnector.from_url(proxyurl) atimeout = aiohttp.ClientTimeout(total=timeout) async with aiohttp.ClientSession(connector=connector, timeout=atimeout) as sess: try: async with sess.request(method, url, headers=headers, params=params, json=json, data=body, ssl=ssl) as resp: info = { 'ok': True, 'url': str(resp.url), 'code': resp.status, 'headers': dict(resp.headers), } hashset = s_hashset.HashSet() async with await self.upload() as upload: async for byts in resp.content.iter_chunked( CHUNK_SIZE): await upload.write(byts) hashset.update(byts) size, _ = await upload.save() info['size'] = size info['hashes'] = dict([(n, s_common.ehex(h)) for (n, h) in hashset.digests()]) return info except asyncio.CancelledError: raise except Exception as e: exc = s_common.excinfo(e) mesg = exc.get('errmsg') if not mesg: mesg = exc.get('err') return { 'ok': False, 'mesg': mesg, }