class DBPool(object):
def __init__(self, DB_HOST, DB_PORT, DB_USER, DB_PASS, DB_DB,
size=10, error_reporting=False):
self.pool = Queue()
for i in range(size):
con = umysql.Connection()
try:
con.connect(DB_HOST, DB_PORT, DB_USER, DB_PASS, DB_DB)
self.pool.put_nowait(con)
except:
sys.exit("Could not connect to database")
self.error_reporting = error_reporting
def execute(self, sql, params=()):
con = self.pool.get()
res = False
try:
if (params):
res = con.query(sql, params)
else:
res = con.query(sql)
except umysql.SQLError, e:
if self.error_reporting:
print "Code: %s\nMessage: %s\nSQL: %s\n"\
% (e[0], e[1], sql.encode("utf-8"))
finally:
class Worker():
def __init__(self):
self.threads = []
self.queue = Queue()
def long_func(self, th, seed):
k = 0
while k < 10000:
print "LOG: Inside the long function Thread: ", th, " Seed: ", seed
time.sleep(.1)
print "LOG: Long function is out of the loop", seed
self.queue.put_nowait(seed)
def short_func(self, th, seed):
print "LOG: Inside the short function Thread:", th, " Seed: ", seed
self.queue.put_nowait(seed)
def start(self, seed):
print "INFO: Initializing the threads..."
self.threads.append(gevent.spawn(self.long_func, 1, seed))
gevent.sleep(1)
self.threads.append(gevent.spawn(self.short_func, 2, seed))
while self.queue.empty():
print "INFO: Queue is empty %s" % seed
gevent.sleep(0)
raise TaskComplete
def stop(self):
gevent.killall(self.threads)
def batch_request(url, reqs, config={}, workers_num=1):
if "interval_time" not in config:
config["interval_time"] = None
if "headers" not in config:
config["headers"] = {}
elif "content-type" in config["headers"]:
if config["headers"]["content-type"] == "application/json":
reqs = [json.dumps(req).strip() for req in reqs]
if "method" not in config:
config["method"] = "POST"
if "expected" not in config:
config["expected"] = "True"
reqs_queue = Queue()
for req in reqs:
reqs_queue.put_nowait(req)
count = reqs_queue.qsize()
workers = [gevent.spawn(request_worker, url, reqs_queue, config, worker=worker) for worker in range(workers_num)]
start = time()
# send requests at the same time
gevent.joinall(workers)
# Response.elapsed for a single request, offered by requests
print "total requests number: %d" % (count)
print "total elapsed time: %s\n" % (time()-start)
class GameServer:
def __init__(self, listen, workers=1):
#initialize the exit channel
self.exitChan = Queue(maxsize=1)
#server is the grpc server
self.server = grpc.server(
futures.ThreadPoolExecutor(max_workers=workers))
#initialize the game
self.game = GoGame(maxsize=10)
#the server here is the grpc server
self.server.add_insecure_port(listen)
#regist the game to the server
#the game here is just the gogame created above
game_pb2_grpc.add_GameServicer_to_server(self.game, self.server)
def start(self):
#start the server
self.server.start()
while self.exitChan.empty():
gevent.sleep(0)
#stop the console of the game
self.game.stop_console()
def stop(self):
#stop the server
self.server.stop(0)
#put a element into the channel , so the channel is not empty and the loop will end
self.exitChan.put_nowait(0)
class Transport(object):
def __init__(self, tid, name):
self.__tid = tid
self.__name = name
self.q_maxsize = 100
self.__task_queue = Queue(maxsize=self.q_maxsize)
def add_task(self, task):
if self.__task_queue.qsize() < self.q_maxsize:
self.__task_queue.put_nowait(task)
return True
else:
# queue is full
logger.debug('The task queue is full')
return False
def get(self, timeout=None):
task = None
try:
task = self.__task_queue.get(block=True, timeout=timeout)
except Exception as e:
if not self.__task_queue.empty():
logger.error('Transport Thread still in processing data')
finally:
return task
@property
def name(self):
return str(self.__name)
class ConnectionPool:
def __init__(self, db_config, time_to_sleep=30, test_run=False):
self.username = db_config.get('user')
self.password = db_config.get('password')
self.host = db_config.get('host')
self.port = int(db_config.get('port'))
self.max_pool_size = 20
self.test_run = test_run
self.pool = None
self.time_to_sleep = time_to_sleep
self._initialize_pool()
def get_initialized_connection_pool(self):
return self.pool
def _initialize_pool(self):
self.pool = Queue(maxsize=self.max_pool_size)
current_pool_size = self.pool.qsize()
if current_pool_size < self.max_pool_size: # this is a redundant check, can be removed
for _ in xrange(0, self.max_pool_size - current_pool_size):
try:
conn = db.connect(host=self.host,
user=self.username,
passwd=self.password,
port=self.port)
self.pool.put_nowait(conn)
except db.OperationalError, e:
LOGGER.error("Cannot initialize connection pool - retrying in {} seconds".format(self.time_to_sleep))
LOGGER.exception(e)
break
self._check_for_connection_loss()
class Qoorate(BrooklynCodeBrubeck):
"""Custom application class for Qoorate."""
def __init__(self, settings_file=None, project_dir=None,
*args, **kwargs):
""" Most of the parameters are dealt with by Brubeck,
Additional functionality follow call to super
"""
super(Qoorate, self).__init__(settings_file, project_dir, **kwargs)
pool_size = 10
if self.db_conn == None:
# create our db_conn if we have the settings to
if settings_file != None:
mysql_settings = self.get_settings('mysql')
if mysql_settings != None:
logging.debug("creating application db_conn pool")
self.db_conn = Queue()
for i in range(pool_size):
self.db_conn.put_nowait(create_db_conn(mysql_settings))
self.template_env.filters['sc'] = unsanitize_safe_htmlentities
def determine_relevency(self, item):
"""schedule an indexing using concurrency"""
logging.info("qoorate_determine_relevency, start: %s" % item)
logging.info("qoorate_generate_relevency, star greenlet: %s" % item)
qoorate_determine_relevency(item)
logging.info("qoorate_generate_relevency, end: %s" % item)
class AEntitySink(Greenlet, IASyncConsumer):
"""
Abstract entity sink which is capable of consuming entities via both
push and pull-based mechanisms.
"""
def __init__(self, desc, maxQueueSize=None):
Greenlet.__init__(self)
self._desc = desc
self._input = Queue(maxQueueSize)
def put(self, item, block=True, timeout=None):
"""Inserts an item into this sink's queue"""
self._input.put(item, block, timeout)
def put_nowait(self, item):
"""Inserts an item into this sink's queue only if it would be non-blocking"""
self._input.put_nowait(item)
def _run(self):
"""Subclasses should override to process the pull-based loop in the
context of this sink's Greenlet."""
pass
def processQueue(self, queue, async=True, poolSize=128):
class Emutiv():
"""Emotiv device emulator."""
channels = [
'F3', 'F4', 'P7', 'FC6', 'F7', 'F8', 'T7', 'P8', 'AF4', 'T8', 'AF3',
'O2', 'O1', 'FC5'
]
def __init__(self, displayOutput=False):
"""Constructs an emulator instance."""
self.file = open(fname, 'rb')
self.input = pickle.load(self.file)
self.packets = Queue()
def setup(self):
"""Sets an emulator up."""
print 'started'
while True:
for x in self.input:
x.__recovery__()
self.packets.put_nowait(x)
gevent.sleep(0.0078125)
def dequeue(self):
"""Read another packet from the file."""
try:
return self.packets.get()
except Exception, e:
print e
def test_investigator(self):
"""
Verify the investigator.
"""
with mock.patch('commissaire.transport.ansibleapi.Transport') as _tp:
_tp().get_info.return_value = (
0,
{
'os': 'fedora',
'cpus': 2,
'memory': 11989228,
'space': 487652,
}
)
q = Queue()
client = etcd.Client()
client.get = MagicMock('get')
client.get.return_value = MagicMock(value=self.etcd_host)
client.set = MagicMock('set')
client.set.return_value = self.etcd_host
to_investigate = {
'address': '10.0.0.2',
}
ssh_priv_key = 'dGVzdAo='
q.put_nowait((to_investigate, ssh_priv_key))
investigator(q, client, True)
self.assertEquals(1, client.get.call_count)
self.assertEquals(1, client.set.call_count)
class Spider:
def __init__(self, url='', depth=1):
self.tasks = Queue()
self.tasks.put(url)
self.init_url = url or ''
self.depth = depth or ''
def run(self):
threds = [
gevent.spawn(self.work),
gevent.spawn(self.work),
gevent.spawn(self.work),
gevent.spawn(self.work)
]
gevent.joinall(threds)
def work(self):
while not self.tasks.empty():
page = self.tasks.get()
p = Page(page, '')
p.do_request()
p.parse_content()
hrefs = p.hrefs
for href in hrefs:
self.tasks.put_nowait(href)
def _run(self):
utils.log("[%s] parsing site %s" % (self, self.base))
queue = Queue()
pool = Pool(32)
seed = 'http://www.amazon.com/best-sellers-books-Amazon/zgbs/books/'
parsed = set()
queue.put_nowait((seed, 'seed', 0))
while True:
items = []
while not queue.empty():
item = queue.get_nowait()
if item[0] not in parsed:
items.append(item)
parsed.add(item[0])
if 0 == len(items) and 0 == len(pool):
break
for item in items:
pool.spawn(self._parseResultsPage, queue, item[0], item[1], item[2])
time.sleep(0.01)
pool.join()
self._output.put(StopIteration)
class ConnectionPool:
def __init__(self, db_config, time_to_sleep=30, test_run=False):
self.username = db_config.get('user')
self.password = db_config.get('password')
self.host = db_config.get('host')
self.port = int(db_config.get('port'))
self.max_pool_size = 20
self.test_run = test_run
self.pool = None
self.time_to_sleep = time_to_sleep
self._initialize_pool()
def get_initialized_connection_pool(self):
return self.pool
def _initialize_pool(self):
self.pool = Queue(maxsize=self.max_pool_size)
current_pool_size = self.pool.qsize()
if current_pool_size < self.max_pool_size: # this is a redundant check, can be removed
for _ in xrange(0, self.max_pool_size - current_pool_size):
try:
conn = db.connect(host=self.host,
user=self.username,
passwd=self.password,
port=self.port)
self.pool.put_nowait(conn)
except db.OperationalError, e:
LOGGER.error(
"Cannot initialize connection pool - retrying in {} seconds"
.format(self.time_to_sleep))
LOGGER.exception(e)
break
self._check_for_connection_loss()
class Worker():
def __init__(self,inputdict, timeout, outputmode, validation_func):
self.threads = []
self.queue = Queue()
self.inputdict = inputdict
self.timeout = timeout
self.outputmode = outputmode
self.validation_func = validation_func
def infi(self, th, thm):
k = 0
while k<10000:
print 'I am in INFI ', th, thm
time.sleep(.1)
print "out while infi", thm
self.queue.put_nowait(thm)
def test(self, th, thm):
print "inside test", thm
self.queue.put_nowait(thm)
def start(self, thm):
print "Hii"
self.threads.append(gevent.spawn(self.infi, 1, thm))
self.threads.append(gevent.spawn(self.test, 2, thm))
while self.queue.empty():
print "queue is empty %s" % thm
gevent.sleep(0)
raise TaskComplete
def stop(self):
gevent.killall(self.threads)
class DownloadByGevent(GetTitleUrls):
def __init__(self, download_urls):
super(DownloadByGevent, self).__init__()
self.workQueue = Queue(1000)
self.download_urls = download_urls
def crawler(self, index):
Process_id = 'Process-' + str(index)
while not self.workQueue.empty():
url = self.workQueue.get(timeout=2)
try:
r = self.get_soup(url, timeout=20)
my_write(str(' '.join([str(x) for x in [Process_id, self.workQueue.qsize(), r.status_code, url]])))
except Exception as e:
# print(Process_id, self.workQueue.qsize, url, 'Error: ', e)
my_write(str(' '.join([str(x) for x in [Process_id, self.workQueue.qsize(), r.status_code, url]])))
def boss(self):
for url in self.download_urls:
self.workQueue.put_nowait(url)
def start(self):
start_time = time.time()
gevent.spawn(self.boss).join()
jobs = []
for i in range(10):
jobs.append(gevent.spawn(self.crawler, i))
gevent.joinall(jobs)
print('end')
print("爬虫时间为%s"%time.time() - start_time)
def test_investigator(self):
"""
Verify the investigator.
"""
with mock.patch('commissaire.transport.ansibleapi.Transport') as _tp:
_tp().get_info.return_value = (0, {
'os': 'fedora',
'cpus': 2,
'memory': 11989228,
'space': 487652,
})
q = Queue()
client = etcd.Client()
client.get = MagicMock('get')
client.get.return_value = MagicMock(value=self.etcd_host)
client.set = MagicMock('set')
client.set.return_value = self.etcd_host
to_investigate = {
'address': '10.0.0.2',
}
ssh_priv_key = 'dGVzdAo='
q.put_nowait((to_investigate, ssh_priv_key))
investigator(q, client, True)
self.assertEquals(1, client.get.call_count)
self.assertEquals(1, client.set.call_count)
class Actor(gevent.Greenlet):
def __init__(self):
self.running = False
self.inbox = Queue()
gevent.Greenlet.__init__(self)
def received(self, message):
"""
Define your subclass to handle incoming messages here...
"""
raise NotImplementedError()
def cast(self, message):
""" Send a message to the actor.
If the actor is busy, the message will be enqueued for later
consumption. There is no return value.
>>> a = Actor()
>>> a.received = lambda msg: msg
>>> a.cast("hello")
"""
self.inbox.put_nowait(message)
def _run(self):
self.running = True
while self.running:
message = self.inbox.get()
self.received(message)
class Spider:
def __init__(self, url='', depth=1, threads=4):
self.url = url
self.depth = depth
self.threads = threads
self.tasks = Queue()
self.bucket = []
def run(self):
self.tasks.put(Task(self.url, self.depth))
threds = [
gevent.spawn(self.worker)
for i in range(self.threads)
]
gevent.joinall(threds)
def worker(self, worker_id=''):
while not self.tasks.empty():
task = self.tasks.get()
if task.url in self.bucket:
# here have a bug
continue
self.bucket.append(task.url)
task.run()
for t in task.subtasks:
self.tasks.put_nowait(t)
def recursive_crawl(url):
all_urls = set()
processing_urls = set()
processed_urls = set()
task_queue = Queue()
data_queue = Queue()
def is_processed(url):
return url in processed_urls
def is_processing(url):
return url in processing_urls
def mark_processed(url):
if is_processing(url):
processing_urls.remove(url)
if is_processed(url):
print('Duplicate processed url {}'.format(url))
else:
processed_urls.add(url)
def mark_processing(url):
processing_urls.add(url)
def add_to_all(url):
if url not in all_urls:
print('Record url {}'.format(url))
all_urls.add(url)
mark_processing(url)
task_queue.put_nowait(url)
# Start workers
workers = []
for i in xrange(10):
workers.append(
gevent.spawn(url_worker,
i, task_queue, data_queue)
)
print('workers', len(workers))
while processing_urls:
if data_queue.empty():
gevent.sleep(0)
continue
done_url, hrefs = data_queue.get()
mark_processed(done_url)
for sub_url in hrefs:
add_to_all(sub_url)
if not is_processed(sub_url) and not is_processing(sub_url):
mark_processing(sub_url)
task_queue.put_nowait(sub_url)
print('Processed', len(processed_urls), 'All', len(all_urls))
print('Total latency', demo_helpers.TOTAL_LATENCY)
def do_job(self,jobid,props,*args,**kwargs):
log.info("do job: %s"%(job))
if props:
props = json.loads(props)
else:
props = {}
payload = {"id":jobid,"status":"building"}
r = requests.post('%s/api/job'%(self.options.master_base_url),data=payload)
log.info("update job %s status,result:%s"%(jobid,r))
if props.has_key("steps"):
steps = props["steps"]
for step in steps:
if step.has_key("shell"):
shell = step["shell"]
log.info("shell: %s"%(shell[:256]))
sub = Popen([shell], stdout=PIPE, stderr=PIPE, shell=True)
stdio_q = Queue()
def handle_stdout():
for l in sub.stdout:
stdio_q.put_nowait((0,l))
def handle_stderr():
for l in sub.stderr:
stdio_q.put_nowait((1,l))
def handle_stdio_q():
#stdout 0 stderr 1 extra 2 end 255
current_text_type = None
stdio_list = []
need_flush = False
timeout = None
while 1:
ttype,text = stdio_q.get()
if ttype!=current_text_type and len(stdio_list)>0:
need_flush = True
if len(stdio_list)>50:
need_flush = True
if need_flush:
text2flush = "".join(stdio_list)
payload = {"id":jobid,"text_type":current_text_type,"stdio_text":text2flush}
r = requests.post('%s/api/job'%(self.options.master_base_url),data=payload)
need_flush = False
if ttype==255:
break
current_text_type = ttype
stdio_list.append(text)
glet_stdout = gevent.spawn(handle_stdout)
glet_stderr = gevent.spawn(handle_stderr)
glet_stdio_q = gevent.spawn(handle_stdio_q)
sub.wait()
stdio_q.put_nowait((255,""))
glet_stdout.kill()
glet_stderr.kill()
payload = {"id":jobid,"status":3}#JOB_STATUS_FINISHED
r = requests.post('%s/api/job'%(self.options.master_base_url),data=payload)
log.info("update job %d status,result:%s"%(jobid,r))
class InteractiveShell(Session):
def __init__(self, id, *args, **kw):
super(InteractiveShell, self).__init__(id, *args, **kw)
self.shell = None
self.cmd_id = 0
self.shell_worker = None
self.commands = Queue()
def start(self):
if self.shell_worker is None:
def worker():
while 1:
command = self.commands.get(True)
if command is None:
break
self.cmd_id = command['id']
out = StringIO()
orig = sys.stdout
sys.stdout = out
more = self.shell.runsource(command['source'])
sys.stdout = orig
self.send({'id': self.cmd_id,
'complete': True,
'more': more,
'out': out.getvalue()})
self.shell_worker = gevent.Greenlet(worker)
if not self.shell_worker:
self.shell_worker.start()
def write(self, data):
print data,
def on_open(self):
self.shell = code.InteractiveInterpreter(
{'session': self,
'registry': self.registry,
'manager': self.manager})
self.shell.write = self.write
self.start()
self.send({'id': 0,
'complete': False,
'more': False,
'out': "Python %s on %s\n" % (sys.version, sys.platform)})
def on_message(self, msg):
self.commands.put_nowait(json.loads(msg))
def on_closed(self):
if self.shell_worker:
self.shell_worker.kill()
class MemorySession(Session):
"""
In memory session with a outgoing gevent Queue as the message
store.
"""
def __init__(self, server, session_id=None):
super(MemorySession, self).__init__(server, session_id=session_id)
self.session_id = session_id or str(uuid.uuid4())[:8]
self.server = server
self.queue = Queue()
self.hits = 0
self.heartbeats = 0
self.connected = False
def add_message(self, msg):
self.queue.put_nowait(msg)
def get_messages(self, **kwargs):
timeout = kwargs.get('timeout', None)
self.incr_hits()
if self.queue.empty():
try:
return self.queue.get(**kwargs)
except Empty:
return []
else:
accum = []
try:
while not self.queue.empty():
if timeout:
accum.append(self.queue.get(timeout=timeout))
else:
accum.append(self.queue.get_nowait())
finally:
return accum
def interrupt(self):
"""
A kill event trigged through a client accessible endpoint
Internal expires will not have is_interupted() == True
"""
self.interrupted = True
self.kill()
def kill(self):
self.connected = False
# Expire only once
if not self.expired:
self.expired = True
self.timeout.set()
class hostpocfactory:
def __init__(self, host, port, threads):
self.host = host
self.port = port
self.threads = threads
self.vulns = list()
self.queue = Queue(self.threads)
def loadmodule(self):
requirements = [
'redis', 'socket', 're', 'pymongo', 'sys', 'os', 'paramiko',
'binascii', 'struct', 'string', 'random', 'codecs', 'time',
'select', 'ssl', 'telnetlib', 'hashlib', 'io', 'string', 'urllib'
]
imported_libs = {
lib: importlib.import_module(lib)
for lib in requirements
}
globals().update(imported_libs)
def Consumer(self, pocstr):
while not self.queue.empty():
data = self.queue.get()
sem.acquire()
exec(data)
gevent.sleep(0)
sem.release()
def runpocwithsysname(self, keyword):
try:
poclist = list()
self.loadmodule()
sql = 'SELECT poc from hostexploit WHERE vulname like "%{}%"'.format(
keyword)
res = db().execute(sql)
for item in res:
poclist.append(item['poc'])
self.queue.put_nowait(item['poc'])
mylog('hostexploit', True).log.info(pyfancy().green(
'[+]针对目标:{0}:{1} 加载{2} hostpoc {3}个'.format(
self.host, self.port, keyword, len(poclist))))
threads = [gevent.spawn(self.Consumer, item) for item in poclist]
gevent.joinall(threads)
for vuln in self.vulns:
sqlstr = 'INSERT INTO hostvulnlist (vulnhost, vulnport, vulnname, isvul, payload, proof, exception) VALUE ("{0}", "{1}", "{2}", "{3}", "{4}", "{5}", "{6}")'.format(
vuln['vulnhost'], vuln['vulnport'],
vuln['vulnname'], vuln['isvul'],
pymysql.escape_string(str(vuln['payload'])), vuln['proof'],
pymysql.escape_string(str(vuln['exception'])))
db().execute(sqlstr)
vuln = json.dumps(vuln, indent=4)
mylog('hostexploit').log.debug(pyfancy().magenta(
'[*] {0}'.format(vuln)))
self.vulns = []
except Exception as e:
mylog('hostexploit').log.critical(pyfancy().red(e))
class Executor(BaseInnerExecutor):
def __init__(self, *args, **kwargs):
self.names = {}
self._output_queue = Queue()
super(SshExecutor.Executor, self).__init__(*args, **kwargs)
@staticmethod
def _stream_fd(fd, queue):
for line in iter(fd.readline, b""):
queue.put_nowait((fd, line))
def _consume(self, queue):
while True:
try:
output = queue.get()
except Empty:
continue
# None is explicitly sent to shutdown the consumer
if output is None:
return
fd, line = output
self.update(self.hostname, self.names[fd], line)
def run(self):
_proc = Popen(
["ssh", "-no", "PasswordAuthentication=no"] + self.parent.ssh_opts + [self.hostname] + self.command,
stdout=PIPE, stderr=PIPE
)
self.names = {
_proc.stdout: "stdout",
_proc.stderr: "stderr",
}
out_worker = gevent.spawn(self._stream_fd, _proc.stdout, self._output_queue)
err_worker = gevent.spawn(self._stream_fd, _proc.stderr, self._output_queue)
waiter = gevent.spawn(_proc.wait)
consumer = gevent.spawn(self._consume, self._output_queue)
gevent.joinall([out_worker, err_worker, waiter], timeout=self.timeout)
# If we've made it here and the process hasn't completed we've timed out.
if _proc.poll() is None:
self._output_queue.put_nowait(
(_proc.stderr, "GSH: command timed out after %s second(s).\n" % self.timeout))
_proc.kill()
rc = _proc.wait()
self._output_queue.put_nowait(None)
consumer.join()
return rc
class geventimg(object):
def __init__(self):
self.file_path = 'D:/animeimg/animeimg2/'
self.queue = Queue()
def getanimeimgurlfrommysql(self):
try:
if not os.path.exists(self.file_path):
print ('文件夹',self.file_path,'不存在，重新建立')
os.makedirs(self.file_path)
db.execute('SELECT a.ANIME_ID,a.ANIME_IMAGE FROM anime_home a',None)
records = db.fetchall()
if records:
for r in records:
self.queue.put_nowait(r)
except Exception as e:
pass
else:
pass
finally:
pass
pass
def downloadimg(self,n):
while not self.queue.empty():
try:
animeimg = self.queue.get_nowait()
r = self.url_open('http://donghua.dmzj.com{}'.format(animeimg[1]))
# print(animeimg)
anime_img_name = self.file_path +'animepic_{}.jpg'.format(animeimg[0])
with open(anime_img_name,"wb") as f:
f.write(r.content)
f.flush()
f.close()
gevent.sleep(1)
except Exception as e:
traceback.print_exc()
finally:
pass
def url_open(self,url):
"""
爬取网页
"""
req = request.get(url,5,'donghua.dmzj.com',None)
return req
def main(self):
gevent.spawn(self.getanimeimgurlfrommysql).join()
# gevent.spawn(self.downloadimg).join()
pool = gevent.pool.Pool(5)
threads = []
for r in range(5):
threads.append(pool.spawn(self.downloadimg,r))
gevent.joinall(threads)
class MemorySession(Session):
"""
In memory session with a outgoing gevent Queue as the message
store.
"""
timer = 10.0
def __init__(self, server, session_id=None):
super(MemorySession, self).__init__(server, session_id)
self.session_id = session_id or str(uuid.uuid4())[:8]
self.server = server
self.queue = Queue()
self.hits = 0
self.heartbeats = 0
self.connected = False
def add_message(self, msg):
self.queue.put_nowait(msg)
def get_messages(self, **kwargs):
self.incr_hits()
if self.queue.empty():
try:
return self.queue.get(**kwargs)
except Empty:
return []
else:
accum = []
try:
while not self.queue.empty():
accum.append(self.queue.get_nowait())
finally:
return accum
def interrupt(self):
"""
A kill event trigged through a client accessible endpoint
Internal expires will not have is_interupted() == True
"""
self.interrupted = True
self.kill()
def kill(self):
self.connected = False
# Expire only once
if not self.expired:
self.expired = True
self.timeout.set()
class MatlabConnect(object):
def __init__(self, headset, path=''):
self.headset = headset
self.dataQueue = Queue()
self.isRunning = False
self.sensors = [
'F3', 'FC5', 'AF3', 'F7', 'T7', 'P7', 'O1', 'O2', 'P8', 'T8', 'F8',
'AF4', 'FC6', 'F4'
]
command = os.path.join(path, 'matlab')
os.system('%s -automation -desktop' % command)
def get_sensors_info(self):
"""
Greenlet to get a packet from Emotiv headset.
Append new data to queues where consumers will read from
"""
try:
while self.isRunning:
packet = self.headset.dequeue()
values = [
packet.sensors[name]['value'] for name in self.sensors
]
if self.dataQueue is not None:
self.dataQueue.put_nowait(values)
gevent.sleep(0)
except KeyboardInterrupt:
print('Read stopped')
self.isRunning = False
except Exception as e:
print('Read Error: %s' % e)
self.isRunning = False
finally:
print('Read over')
self.isRunning = False
self.headset.close()
def matlabBridge(self, varName):
data = None
while self.isRunning or not self.dataQueue.empty():
while not self.dataQueue.empty():
buf = self.dataQueue.get()
if data is None:
data = np.array(buf, dtype=int)
else:
data = np.vstack((data, buf))
self.session.putvalue(varName, data)
gevent.sleep(1)
print 'Matlab over'
def url_crawler():
global work
work = Queue() #实例化队列对象
for group in range(2):
for page in range(2):
url='http://www.boohee.com/food/group/%s?page=%s'%(group+1,page+1)
res=requests.get(url,headers=headers)
soup=BeautifulSoup(res.text,'html.parser')
url_sort=soup.find_all(class_='text-box pull-left')
for i in range(len(url_sort)):
food_url='http://www.boohee.com'+url_sort[i].find('a')['href'] #获取食物详情页的链接
work.put_nowait(food_url) #向队列中添加url
return work
class AsyncWorker(Worker):
def __init__(self, worker_nums=5):
self._queue = Queue()
for _ in range(worker_nums):
gevent.spawn(self._start_workers)
def _start_workers(self):
while True:
func, args, kwargs = self._queue.get()
func(*args, **kwargs)
def run(self, func, *args, **kwargs):
self._queue.put_nowait((func, args, kwargs))
class MatlabConnect(object):
def __init__(self, headset, path = ''):
self.headset = headset
self.dataQueue = Queue()
self.isRunning = False
self.sensors = ['F3','FC5', 'AF3', 'F7', 'T7', 'P7', 'O1', 'O2', 'P8', 'T8', 'F8', 'AF4', 'FC6', 'F4']
command = os.path.join(path, 'matlab')
os.system('%s -automation -desktop' % command)
def get_sensors_info(self):
"""
Greenlet to get a packet from Emotiv headset.
Append new data to queues where consumers will read from
"""
try:
while self.isRunning:
packet = self.headset.dequeue()
values = [packet.sensors[name]['value'] for name in self.sensors]
if self.dataQueue is not None:
self.dataQueue.put_nowait(values)
gevent.sleep(0)
except KeyboardInterrupt:
print ('Read stopped')
self.isRunning = False
except Exception as e:
print ('Read Error: %s' % e)
self.isRunning = False
finally:
print ('Read over')
self.isRunning = False
self.headset.close()
def matlabBridge(self, varName):
data = None
while self.isRunning or not self.dataQueue.empty():
while not self.dataQueue.empty():
buf = self.dataQueue.get()
if data is None:
data = np.array(buf, dtype=int)
else:
data = np.vstack((data, buf))
self.session.putvalue(varName, data)
gevent.sleep(1)
print 'Matlab over'
def save_imgs(self, jpgs: 'list', gifs: 'list', title: 'str'):
'''解析图片元素列表，存储图片到本地
:param jpgs: jpg的源地址列表
:param gifs: gif的源地址列表
:param title: 保存的文件夹名称
:return: None
'''
def _save():
while not work.empty():
img = work.get_nowait()
url = img
# 检查是否为gif
if '.gif' in url:
self.gif += 1
pathName = "./%s/动图%d.gif" % (title, self.gif)
print('jpg-已完成%d张，共%d张；gif-已完成%d张，共%d张。' %
(self.jpg, self.jpg_num, self.gif, self.gif_num))
urllib.request.urlretrieve(url, pathName)
else:
self.jpg += 1
pathName = "./%s/图片%d.jpg" % (title, self.jpg)
print('jpg-已完成%d张，共%d张；gif-已完成%d张，共%d张。' %
(self.jpg, self.jpg_num, self.gif, self.gif_num))
urllib.request.urlretrieve(url, pathName)
# 建立队列
work = Queue()
# 建立图片原地址列表
for url in jpgs:
work.put_nowait(url)
for url in gifs:
work.put_nowait(url)
# 添加任务到队列并开始执行
tasks_list = []
for x in range(4):
task = gevent.spawn(_save)
tasks_list.append(task)
gevent.joinall(tasks_list)
print('一共爬取了%d个jpg表情和%d个gif动图' % (self.jpg, self.gif))
class ZMQSummarizedTestResult(ZMQTestResult):
def __init__(self, args):
super(ZMQSummarizedTestResult, self).__init__(args)
self.interval = 1.
self._data = Queue()
gevent.spawn_later(self.interval, self._dump_data)
def push(self, data_type, **data):
self._data.put_nowait((data_type, data))
def close(self):
while not self._data.empty():
self._dump_data(loop=False)
self.context.destroy()
def _dump_data(self, loop=True):
if self._data.empty() and loop:
gevent.spawn_later(self.interval, self._dump_data)
return
data = {
'data_type': 'batch',
'agent_id': self.agent_id,
'hostname': get_hostname(),
'run_id': self.run_id,
'counts': defaultdict(list)
}
# grabbing what we have
for _ in range(self._data.qsize()):
data_type, message = self._data.get()
data['counts'][data_type].append(message)
while True:
try:
self._push.send(self.encoder.encode(data), zmq.NOBLOCK)
break
except zmq.ZMQError as e:
if e.errno in (errno.EAGAIN, errno.EWOULDBLOCK):
continue
else:
raise
if loop:
gevent.spawn_later(self.interval, self._dump_data)
class ZMQSummarizedTestResult(ZMQTestResult):
def __init__(self, args):
super(ZMQSummarizedTestResult, self).__init__(args)
self.interval = 1.
self._data = Queue()
gevent.spawn_later(self.interval, self._dump_data)
def push(self, data_type, **data):
self._data.put_nowait((data_type, data))
def close(self):
while not self._data.empty():
self._dump_data(loop=False)
self.context.destroy()
def _dump_data(self, loop=True):
if self._data.empty() and loop:
gevent.spawn_later(self.interval, self._dump_data)
return
data = {'data_type': 'batch',
'agent_id': self.agent_id,
'hostname': get_hostname(),
'run_id': self.run_id,
'counts': defaultdict(list)}
# grabbing what we have
for _ in range(self._data.qsize()):
data_type, message = self._data.get()
data['counts'][data_type].append(message)
while True:
try:
self._push.send(self.encoder.encode(data), zmq.NOBLOCK)
break
except zmq.ZMQError as e:
if e.errno in (errno.EAGAIN, errno.EWOULDBLOCK):
continue
else:
raise
if loop:
gevent.spawn_later(self.interval, self._dump_data)
def test_investigator(self):
"""
Verify the investigator.
"""
with mock.patch('commissaire.transport.ansibleapi.Transport') as _tp:
_tp().get_info.return_value = (
0,
{
'os': 'fedora',
'cpus': 2,
'memory': 11989228,
'space': 487652,
}
)
q = Queue()
client = etcd.Client()
client.get = MagicMock('get')
client.get.return_value = MagicMock(value=self.etcd_host)
client.set = MagicMock('set')
client.set.return_value = self.etcd_host
to_investigate = {
'address': '10.0.0.2',
}
ssh_priv_key = 'dGVzdAo='
connection_config = {
'etcd': {
'uri': urlparse('http://127.0.0.1:2379'),
},
'kubernetes': {
'uri': urlparse('http://127.0.0.1:8080'),
'token': 'token',
}
}
q.put_nowait((to_investigate, ssh_priv_key))
investigator(q, connection_config, client, True)
self.assertEquals(1, client.get.call_count)
self.assertEquals(2, client.set.call_count)
def download_to_s3(self, start_date, end_date):
days = list(daterange(start_date, end_date))
logger.info("Processing {} days of data".format(len(days)))
for day in days:
logger.info("Adding %d reports to the queue" % len(self.api_keys))
tasks = Queue()
for x in self.api_keys:
if x.api_key:
tasks.put_nowait((x, day.strftime(DATE_FORMAT)))
# This statement will not continue until all workers are done/cancelled/killed
gevent.joinall([
gevent.spawn(self.run_async_report_worker, i, tasks)
for i in xrange(self.max_parallel_reports)
])
logger.info("Finished all queries for {}".format(
day.strftime(DATE_FORMAT)))
class MyCrawler(object):
def __init__(self):
self.pool = Pool(10)
self.tasks = Queue()
def init_tasks(self):
resp = requests.get('http://www.ttmeiju.com/summary.html')
assert resp.status_code==200
soup = BeautifulSoup(resp.text)
tb = soup.find('table')
def mya(tag):
return tag.name=='a' and tag.parent.name=='td'
meiju={}
for atag in tb.find_all(mya)[3:]:
meiju[atag.string] = atag['href']
self.tasks.put_nowait({'key':atag.string, 'val': atag['href']})
print (meiju)
def do_jobs(self):
self.init_tasks()
while not self.tasks.empty():
task = self.tasks.get()
self.before(task)
self.pool.spawn(self.work, task).rawlink(self.after)
# 等待任务完成
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
self.pool.join()
def work(self, task):
print ('子类实现')
def before(self, task):
print('')
def after(self, asynret):
print('')
class AsyncSendSocket(socket):
def __init__(self, *args, **kwargs):
socket.__init__(self, *args, **kwargs)
self._send_queue = Queue()
self._send_thread = Greenlet.spawn(self._run_async_send)
def _run_async_send(self):
_queue = self._send_queue
while True:
msg = _queue.get()
if msg == _close_message:
self.close()
break
else:
self.sendall(msg)
def async_send(self, msg):
#print 'async send', repr(msg)
self._send_queue.put(msg)
def async_close(self):
self._send_queue.put_nowait(_close_message)
def _concurrent_execute(self, context, start_req, parser, pool, pool_size):
queue = Queue() # 任务队列
# 将初始化请求加入任务队列
for r in start_req:
queue.put_nowait(r)
if pool is None:
pool = GeventPool(pool_size)
greenlets = []
while True:
try:
req = self._check_req(queue.get(timeout=1))
if req.parser is None:
req.parser = parser
greenlets.append(pool.spawn(req, context, queue))
except Empty:
break
return [greenlet.get() for greenlet in greenlets]
def _concurrent_execute(self, context, start_req, parser, pool, pool_size):
queue = Queue() # 任务队列
# 将初始化请求加入任务队列
for r in start_req:
queue.put_nowait(r)
if pool is None:
pool = GeventPool(pool_size)
greenlets = []
while True:
try:
req = self._check_req(queue.get(timeout=1))
if req.parser is None:
req.parser = parser
greenlets.append(pool.spawn(req, context, queue))
except Empty:
break
return [greenlet.get() for greenlet in greenlets]
def test_investigator(self):
"""
Verify the investigator.
"""
with mock.patch('commissaire.transport.ansibleapi.Transport') as _tp:
_tp().get_info.return_value = (0, {
'os': 'fedora',
'cpus': 2,
'memory': 11989228,
'space': 487652,
})
q = Queue()
client = etcd.Client()
client.get = MagicMock('get')
client.get.return_value = MagicMock(value=self.etcd_host)
client.set = MagicMock('set')
client.set.return_value = self.etcd_host
to_investigate = {
'address': '10.0.0.2',
}
ssh_priv_key = 'dGVzdAo='
connection_config = {
'etcd': {
'uri': urlparse('http://127.0.0.1:2379'),
},
'kubernetes': {
'uri': urlparse('http://127.0.0.1:8080'),
'token': 'token',
}
}
q.put_nowait((to_investigate, ssh_priv_key))
investigator(q, connection_config, client, True)
self.assertEquals(1, client.get.call_count)
self.assertEquals(2, client.set.call_count)
class MemorySession(session.Session):
"""
In memory session with a ``gevent.pool.Queue`` as the message store.
"""
def __init__(self, *args, **kwargs):
super(MemorySession, self).__init__(*args, **kwargs)
self._queue = Queue()
def add_message(self, frame, data):
log.info('session closed: %s', self.id)
self._queue.put_nowait((frame, data))
#waiter = self._waiter
#if waiter is not None:
# self._waiter = None
# if not waiter.cancelled():
# waiter.set_result(True)
self._tick()
def get_messages(self, timeout=None):
self._tick()
messages = []
# get all messages immediately pending in the queue
while not self._queue.empty():
try:
msg = self._queue.get_nowait()
except Empty:
break
messages.append(msg)
if not messages:
try:
messages.append(self._queue.get(timeout=timeout))
except Empty:
pass
return messages
class RequestIterator(object):
def __init__(self):
self.queue = Queue()
def __iter__(self):
return self
def _next(self):
data = self.queue.get(True)
if data[0] == stock_provider_pb2.SimulationMsgType.MSG_TICK:
return stock_provider_pb2.SimulationMsg(msgtype=data[0],
tick=data[1])
elif data[0] == stock_provider_pb2.SimulationMsgType.MSG_BIDASK:
return stock_provider_pb2.SimulationMsg(msgtype=data[0],
bidask=data[1])
elif data[0] == stock_provider_pb2.SimulationMsgType.MSG_SUBJECT:
return stock_provider_pb2.SimulationMsg(msgtype=data[0],
subject=data[1])
elif data[0] == stock_provider_pb2.SimulationMsgType.MSG_ALARM:
return stock_provider_pb2.SimulationMsg(msgtype=data[0],
alarm=data[1])
print('Unknown Message')
return stock_provider_pb2.SimulationMsg()
def __next__(self):
return self._next()
def next(self):
return self._next()
def append_tick(self, tick):
self.queue.put_nowait(
(stock_provider_pb2.SimulationMsgType.MSG_TICK, tick))
def append_subject(self, subject):
self.queue.put_nowait(
(stock_provider_pb2.SimulationMsgType.MSG_SUBJECT, subject))
def append_bidask(self, bidask):
self.queue.put_nowait(
(stock_provider_pb2.SimulationMsgType.MSG_BIDASK, bidask))
def append_alarm(self, alarm):
self.queue.put_nowait(
(stock_provider_pb2.SimulationMsgType.MSG_ALARM, alarm))
class MemoryScheduler(Scheduler):
def __init__(self, scheduler_name):
self._scheduler_key = 'Scheduler:{scheduler_name}'.format(
scheduler_name=scheduler_name)
self._queue = Queue()
@classmethod
def from_spider(cls, spider):
scheduler_name = spider.settings.get('SPIDER_NAME')
return cls(scheduler_name)
def __len__(self):
return self._queue.qsize()
def push(self, r):
self._queue.put_nowait(r)
def pull(self,
count=DEFAULT_SCHEDULER_PULL_COUNT,
timeout=DEFAULT_SCHEDULER_PULL_TIMEOUT):
try:
return self._queue.get(timeout=timeout)
except Empty:
raise SchedulerEmpty
class Connection(object):
in_buffer_size = 4096
out_buffer_size = 4096
cql_version = None
protocol_version = 2
keyspace = None
compression = True
compressor = None
decompressor = None
ssl_options = None
last_error = None
in_flight = 0
is_defunct = False
is_closed = False
lock = None
is_control_connection = False
def __init__(self, host='127.0.0.1', port=9042, authenticator=None,
ssl_options=None, sockopts=None, compression=True,
cql_version=None, protocol_version=2, is_control_connection=False):
self.host = host
self.port = port
self.authenticator = authenticator
self.ssl_options = ssl_options
self.sockopts = sockopts
self.compression = compression
self.cql_version = cql_version
self.protocol_version = protocol_version
self.is_control_connection = is_control_connection
self._push_watchers = defaultdict(set)
self._id_queue = Queue(MAX_STREAM_PER_CONNECTION)
for i in range(MAX_STREAM_PER_CONNECTION):
self._id_queue.put_nowait(i)
self.lock = RLock()
@classmethod
def initialize_reactor(self):
"""
Called once by Cluster.connect(). This should be used by implementations
to set up any resources that will be shared across connections.
"""
pass
def close(self):
raise NotImplementedError()
def defunct(self, exc):
with self.lock:
if self.is_defunct or self.is_closed:
return
self.is_defunct = True
log.debug("Defuncting connection (%s) to %s:",
id(self), self.host, exc_info=exc)
self.last_error = exc
self.close()
self.error_all_callbacks(exc)
self.connected_event.set()
return exc
def error_all_callbacks(self, exc):
with self.lock:
callbacks = self._callbacks
self._callbacks = {}
new_exc = ConnectionShutdown(str(exc))
for cb in callbacks.values():
try:
cb(new_exc)
except Exception:
log.warning("Ignoring unhandled exception while erroring callbacks for a "
"failed connection (%s) to host %s:",
id(self), self.host, exc_info=True)
def handle_pushed(self, response):
log.debug("Message pushed from server: %r", response)
for cb in self._push_watchers.get(response.event_type, []):
try:
cb(response.event_args)
except Exception:
log.exception("Pushed event handler errored, ignoring:")
def send_msg(self, msg, cb, wait_for_id=False):
if self.is_defunct:
raise ConnectionShutdown("Connection to %s is defunct" % self.host)
elif self.is_closed:
raise ConnectionShutdown("Connection to %s is closed" % self.host)
if not wait_for_id:
try:
request_id = self._id_queue.get_nowait()
except Empty:
raise ConnectionBusy(
"Connection to %s is at the max number of requests" % self.host)
else:
request_id = self._id_queue.get()
self._callbacks[request_id] = cb
self.push(msg.to_binary(request_id, self.protocol_version, compression=self.compressor))
return request_id
def wait_for_response(self, msg, timeout=None):
return self.wait_for_responses(msg, timeout=timeout)[0]
def wait_for_responses(self, *msgs, **kwargs):
if self.is_closed or self.is_defunct:
raise ConnectionShutdown("Connection %s is already closed" % (self, ))
timeout = kwargs.get('timeout')
waiter = ResponseWaiter(self, len(msgs))
# busy wait for sufficient space on the connection
messages_sent = 0
while True:
needed = len(msgs) - messages_sent
with self.lock:
available = min(needed, MAX_STREAM_PER_CONNECTION - self.in_flight)
self.in_flight += available
for i in range(messages_sent, messages_sent + available):
self.send_msg(msgs[i], partial(waiter.got_response, index=i), wait_for_id=True)
messages_sent += available
if messages_sent == len(msgs):
break
else:
if timeout is not None:
timeout -= 0.01
if timeout <= 0.0:
raise OperationTimedOut()
time.sleep(0.01)
try:
return waiter.deliver(timeout)
except OperationTimedOut:
raise
except Exception as exc:
self.defunct(exc)
raise
def register_watcher(self, event_type, callback):
raise NotImplementedError()
def register_watchers(self, type_callback_dict):
raise NotImplementedError()
def control_conn_disposed(self):
self.is_control_connection = False
self._push_watchers = {}
@defunct_on_error
def process_msg(self, msg, body_len):
version, flags, stream_id, opcode = header_unpack(msg[:4])
if stream_id < 0:
callback = None
else:
callback = self._callbacks.pop(stream_id, None)
self._id_queue.put_nowait(stream_id)
body = None
try:
# check that the protocol version is supported
given_version = version & PROTOCOL_VERSION_MASK
if given_version != self.protocol_version:
msg = "Server protocol version (%d) does not match the specified driver protocol version (%d). " +\
"Consider setting Cluster.protocol_version to %d."
raise ProtocolError(msg % (given_version, self.protocol_version, given_version))
# check that the header direction is correct
if version & HEADER_DIRECTION_MASK != HEADER_DIRECTION_TO_CLIENT:
raise ProtocolError(
"Header direction in response is incorrect; opcode %04x, stream id %r"
% (opcode, stream_id))
if body_len > 0:
body = msg[8:]
elif body_len == 0:
body = six.binary_type()
else:
raise ProtocolError("Got negative body length: %r" % body_len)
response = decode_response(stream_id, flags, opcode, body, self.decompressor)
except Exception as exc:
log.exception("Error decoding response from Cassandra. "
"opcode: %04x; message contents: %r", opcode, msg)
if callback is not None:
callback(exc)
self.defunct(exc)
return
try:
if stream_id < 0:
self.handle_pushed(response)
elif callback is not None:
callback(response)
except Exception:
log.exception("Callback handler errored, ignoring:")
@defunct_on_error
def _send_options_message(self):
if self.cql_version is None and (not self.compression or not locally_supported_compressions):
log.debug("Not sending options message for new connection(%s) to %s "
"because compression is disabled and a cql version was not "
"specified", id(self), self.host)
self._compressor = None
self.cql_version = DEFAULT_CQL_VERSION
self._send_startup_message()
else:
log.debug("Sending initial options message for new connection (%s) to %s", id(self), self.host)
self.send_msg(OptionsMessage(), self._handle_options_response)
@defunct_on_error
def _handle_options_response(self, options_response):
if self.is_defunct:
return
if not isinstance(options_response, SupportedMessage):
if isinstance(options_response, ConnectionException):
raise options_response
else:
log.error("Did not get expected SupportedMessage response; " \
"instead, got: %s", options_response)
raise ConnectionException("Did not get expected SupportedMessage " \
"response; instead, got: %s" \
% (options_response,))
log.debug("Received options response on new connection (%s) from %s",
id(self), self.host)
supported_cql_versions = options_response.cql_versions
remote_supported_compressions = options_response.options['COMPRESSION']
if self.cql_version:
if self.cql_version not in supported_cql_versions:
raise ProtocolError(
"cql_version %r is not supported by remote (w/ native "
"protocol). Supported versions: %r"
% (self.cql_version, supported_cql_versions))
else:
self.cql_version = supported_cql_versions[0]
self._compressor = None
compression_type = None
if self.compression:
overlap = (set(locally_supported_compressions.keys()) &
set(remote_supported_compressions))
if len(overlap) == 0:
log.debug("No available compression types supported on both ends."
" locally supported: %r. remotely supported: %r",
locally_supported_compressions.keys(),
remote_supported_compressions)
else:
compression_type = None
if isinstance(self.compression, six.string_types):
# the user picked a specific compression type ('snappy' or 'lz4')
if self.compression not in remote_supported_compressions:
raise ProtocolError(
"The requested compression type (%s) is not supported by the Cassandra server at %s"
% (self.compression, self.host))
compression_type = self.compression
else:
# our locally supported compressions are ordered to prefer
# lz4, if available
for k in locally_supported_compressions.keys():
if k in overlap:
compression_type = k
break
# set the decompressor here, but set the compressor only after
# a successful Ready message
self._compressor, self.decompressor = \
locally_supported_compressions[compression_type]
self._send_startup_message(compression_type)
@defunct_on_error
def _send_startup_message(self, compression=None):
opts = {}
if compression:
opts['COMPRESSION'] = compression
sm = StartupMessage(cqlversion=self.cql_version, options=opts)
self.send_msg(sm, cb=self._handle_startup_response)
@defunct_on_error
def _handle_startup_response(self, startup_response, did_authenticate=False):
if self.is_defunct:
return
if isinstance(startup_response, ReadyMessage):
log.debug("Got ReadyMessage on new connection (%s) from %s", id(self), self.host)
if self._compressor:
self.compressor = self._compressor
self.connected_event.set()
elif isinstance(startup_response, AuthenticateMessage):
log.debug("Got AuthenticateMessage on new connection (%s) from %s: %s",
id(self), self.host, startup_response.authenticator)
if self.authenticator is None:
raise AuthenticationFailed('Remote end requires authentication.')
self.authenticator_class = startup_response.authenticator
if isinstance(self.authenticator, dict):
log.debug("Sending credentials-based auth response on %s", self)
cm = CredentialsMessage(creds=self.authenticator)
callback = partial(self._handle_startup_response, did_authenticate=True)
self.send_msg(cm, cb=callback)
else:
log.debug("Sending SASL-based auth response on %s", self)
initial_response = self.authenticator.initial_response()
initial_response = "" if initial_response is None else initial_response.encode('utf-8')
self.send_msg(AuthResponseMessage(initial_response), self._handle_auth_response)
elif isinstance(startup_response, ErrorMessage):
log.debug("Received ErrorMessage on new connection (%s) from %s: %s",
id(self), self.host, startup_response.summary_msg())
if did_authenticate:
raise AuthenticationFailed(
"Failed to authenticate to %s: %s" %
(self.host, startup_response.summary_msg()))
else:
raise ConnectionException(
"Failed to initialize new connection to %s: %s"
% (self.host, startup_response.summary_msg()))
elif isinstance(startup_response, ConnectionShutdown):
log.debug("Connection to %s was closed during the startup handshake", (self.host))
raise startup_response
else:
msg = "Unexpected response during Connection setup: %r"
log.error(msg, startup_response)
raise ProtocolError(msg % (startup_response,))
@defunct_on_error
def _handle_auth_response(self, auth_response):
if self.is_defunct:
return
if isinstance(auth_response, AuthSuccessMessage):
log.debug("Connection %s successfully authenticated", self)
self.authenticator.on_authentication_success(auth_response.token)
if self._compressor:
self.compressor = self._compressor
self.connected_event.set()
elif isinstance(auth_response, AuthChallengeMessage):
response = self.authenticator.evaluate_challenge(auth_response.challenge)
msg = AuthResponseMessage("" if response is None else response)
log.debug("Responding to auth challenge on %s", self)
self.send_msg(msg, self._handle_auth_response)
elif isinstance(auth_response, ErrorMessage):
log.debug("Received ErrorMessage on new connection (%s) from %s: %s",
id(self), self.host, auth_response.summary_msg())
raise AuthenticationFailed(
"Failed to authenticate to %s: %s" %
(self.host, auth_response.summary_msg()))
elif isinstance(auth_response, ConnectionShutdown):
log.debug("Connection to %s was closed during the authentication process", self.host)
raise auth_response
else:
msg = "Unexpected response during Connection authentication to %s: %r"
log.error(msg, self.host, auth_response)
raise ProtocolError(msg % (self.host, auth_response))
def set_keyspace_blocking(self, keyspace):
if not keyspace or keyspace == self.keyspace:
return
query = QueryMessage(query='USE "%s"' % (keyspace,),
consistency_level=ConsistencyLevel.ONE)
try:
result = self.wait_for_response(query)
except InvalidRequestException as ire:
# the keyspace probably doesn't exist
raise ire.to_exception()
except Exception as exc:
conn_exc = ConnectionException(
"Problem while setting keyspace: %r" % (exc,), self.host)
self.defunct(conn_exc)
raise conn_exc
if isinstance(result, ResultMessage):
self.keyspace = keyspace
else:
conn_exc = ConnectionException(
"Problem while setting keyspace: %r" % (result,), self.host)
self.defunct(conn_exc)
raise conn_exc
def set_keyspace_async(self, keyspace, callback):
"""
Use this in order to avoid deadlocking the event loop thread.
When the operation completes, `callback` will be called with
two arguments: this connection and an Exception if an error
occurred, otherwise :const:`None`.
"""
if not keyspace or keyspace == self.keyspace:
callback(self, None)
return
query = QueryMessage(query='USE "%s"' % (keyspace,),
consistency_level=ConsistencyLevel.ONE)
def process_result(result):
if isinstance(result, ResultMessage):
self.keyspace = keyspace
callback(self, None)
elif isinstance(result, InvalidRequestException):
callback(self, result.to_exception())
else:
callback(self, self.defunct(ConnectionException(
"Problem while setting keyspace: %r" % (result,), self.host)))
self.send_msg(query, process_result, wait_for_id=True)
def __str__(self):
status = ""
if self.is_defunct:
status = " (defunct)"
elif self.is_closed:
status = " (closed)"
return "<%s(%r) %s:%d%s>" % (self.__class__.__name__, id(self), self.host, self.port, status)
__repr__ = __str__
class Session(object):
"""
Client session which checks the connection health and the queues for
message passing.
"""
def __init__(self):
self.session_id = str(random.random())[2:]
self.client_queue = Queue() # queue for messages to client
self.server_queue = Queue() # queue for messages to server
self.hits = 0
self.heartbeats = 0
self.connected = False
self.timeout = Event()
self.wsgi_app_greenlet = None
def disconnect_timeout():
self.timeout.clear()
if self.timeout.wait(10.0):
gevent.spawn(disconnect_timeout)
else:
self.kill()
gevent.spawn(disconnect_timeout)
def __str__(self):
result = ['session_id=%r' % self.session_id]
if self.connected:
result.append('connected')
if self.client_queue.qsize():
result.append('client_queue[%s]' % self.client_queue.qsize())
if self.server_queue.qsize():
result.append('server_queue[%s]' % self.server_queue.qsize())
if self.hits:
result.append('hits=%s' % self.hits)
if self.heartbeats:
result.append('heartbeats=%s' % self.heartbeats)
return ' '.join(result)
def incr_hits(self):
self.hits += 1
def clear_disconnect_timeout(self):
self.timeout.set()
def heartbeat(self):
self.heartbeats += 1
return self.heartbeats
def valid_heartbeat(self, counter):
self.clear_disconnect_timeout()
return self.heartbeats == counter
def is_new(self):
return self.hits == 0
def kill(self):
if self.connected:
self.connected = False
self.server_queue.put_nowait(None)
self.client_queue.put_nowait(None)
else:
pass # Fail silently
def put_server_msg(self, msg):
self.clear_disconnect_timeout()
self.server_queue.put_nowait(msg)
def put_client_msg(self, msg):
self.clear_disconnect_timeout()
self.client_queue.put_nowait(msg)
def get_client_msg(self, **kwargs):
return self.client_queue.get(**kwargs)
def get_server_msg(self, **kwargs):
return self.server_queue.get(**kwargs)
class Emotiv(object):
"""
Receives, decrypts and stores packets received from Emotiv Headsets.
"""
def __init__(self, display_output=False, serial_number="", is_research=False):
"""
Sets up initial values.
"""
self.running = True
self.packets = Queue()
self.packets_received = 0
self.packets_processed = 0
self.battery = 0
self.display_output = display_output
self.is_research = is_research
self.sensors = {
'F3': {'value': 0, 'quality': 0},
'FC6': {'value': 0, 'quality': 0},
'P7': {'value': 0, 'quality': 0},
'T8': {'value': 0, 'quality': 0},
'F7': {'value': 0, 'quality': 0},
'F8': {'value': 0, 'quality': 0},
'T7': {'value': 0, 'quality': 0},
'P8': {'value': 0, 'quality': 0},
'AF4': {'value': 0, 'quality': 0},
'F4': {'value': 0, 'quality': 0},
'AF3': {'value': 0, 'quality': 0},
'O2': {'value': 0, 'quality': 0},
'O1': {'value': 0, 'quality': 0},
'FC5': {'value': 0, 'quality': 0},
'X': {'value': 0, 'quality': 0},
'Y': {'value': 0, 'quality': 0},
'Unknown': {'value': 0, 'quality': 0}
}
self.serial_number = serial_number # You will need to set this manually for OS X.
self.old_model = False
def setup(self):
"""
Runs setup function depending on platform.
"""
print system_platform + " detected."
if system_platform == "Windows":
self.setup_windows()
elif system_platform == "Linux":
self.setup_posix()
elif system_platform == "Darwin":
self.setup_darwin()
def setup_windows(self):
"""
Setup for headset on the Windows platform.
"""
devices = []
try:
for device in hid.find_all_hid_devices():
if device.vendor_id != 0x21A1 and device.vendor_id != 0xED02:
continue
if device.product_name == 'Brain Waves':
devices.append(device)
device.open()
self.serial_number = device.serial_number
device.set_raw_data_handler(self.handler)
elif device.product_name == 'EPOC BCI':
devices.append(device)
device.open()
self.serial_number = device.serial_number
device.set_raw_data_handler(self.handler)
elif device.product_name == '00000000000':
devices.append(device)
device.open()
self.serial_number = device.serial_number
device.set_raw_data_handler(self.handler)
elif device.product_name == 'Emotiv RAW DATA':
devices.append(device)
device.open()
self.serial_number = device.serial_number
device.set_raw_data_handler(self.handler)
crypto = gevent.spawn(self.setup_crypto, self.serial_number)
console_updater = gevent.spawn(self.update_console)
while self.running:
try:
gevent.sleep(0)
except KeyboardInterrupt:
self.running = False
finally:
for device in devices:
device.close()
gevent.kill(crypto, KeyboardInterrupt)
gevent.kill(console_updater, KeyboardInterrupt)
def handler(self, data):
"""
Receives packets from headset for Windows. Sends them to a Queue to be processed
by the crypto greenlet.
"""
assert data[0] == 0
tasks.put_nowait(''.join(map(chr, data[1:])))
self.packets_received += 1
return True
def setup_posix(self):
"""
Setup for headset on the Linux platform.
Receives packets from headset and sends them to a Queue to be processed
by the crypto greenlet.
"""
_os_decryption = False
if os.path.exists('/dev/eeg/raw'):
# The decryption is handled by the Linux epoc daemon. We don't need to handle it.
_os_decryption = True
hidraw = open("/dev/eeg/raw")
else:
serial, hidraw_filename = get_linux_setup()
self.serial_number = serial
if os.path.exists("/dev/" + hidraw_filename):
hidraw = open("/dev/" + hidraw_filename)
else:
hidraw = open("/dev/hidraw4")
crypto = gevent.spawn(self.setup_crypto, self.serial_number)
console_updater = gevent.spawn(self.update_console)
while self.running:
try:
data = hidraw.read(32)
if data != "":
if _os_decryption:
self.packets.put_nowait(EmotivPacket(data))
else:
#Queue it!
self.packets_received += 1
tasks.put_nowait(data)
gevent.sleep(0)
else:
# No new data from the device; yield
# We cannot sleep(0) here because that would go 100% CPU if both queues are empty
gevent.sleep(DEVICE_POLL_INTERVAL)
except KeyboardInterrupt:
self.running = False
hidraw.close()
if not _os_decryption:
gevent.kill(crypto, KeyboardInterrupt)
gevent.kill(console_updater, KeyboardInterrupt)
def setup_darwin(self):
"""
Setup for headset on the OS X platform.
Receives packets from headset and sends them to a Queue to be processed
by the crypto greenlet.
"""
# Set this to True if the OS is performing the encryption of the packets
_os_decryption = False
# Change these values to the hex equivalent from the output of hid_enumerate. If they are incorrect.
# Current values = VendorID: 8609 ProductID: 1
#hidraw = hid.device(0x31a1, 0x2001)
hidraw = hid.device(0x1234,0xed02)
hidraw.open(0x1234, 0xed02)
self.serial_number = 'SN20120229000290'
if not hidraw:
hidraw = hid.device(0x21a1, 0x1234)
if not hidraw:
hidraw = hid.device(0xed02, 0x1234)
if not hidraw:
print "Device not found. Uncomment the code in setup_darwin and modify hid.device(vendor_id, product_id)"
raise ValueError
if self.serial_number == "":
print "Serial number needs to be specified manually in __init__()."
raise ValueError
print "Serial number:" + self.serial_number
crypto = gevent.spawn(self.setup_crypto, self.serial_number)
console_updater = gevent.spawn(self.update_console)
zero = 0
while self.running:
try:
# Doesn't seem to matter how big we make the buffer 32 returned every time, 33 for other platforms
data = hidraw.read(34,10)
#data = [48]*32
if len(data) == 32:
# Most of the time the 0 is truncated? That's ok we'll add it...
data = [zero] + data
if data != "":
if _os_decryption:
self.packets.put_nowait(EmotivPacket(data))
else:
#Queue it!
print ('Queuing package:'+len(data))
tasks.put_nowait(''.join(map(chr, data[1:])))
self.packets_received += 1
print ('Waiting...')
gevent.sleep(0.01)
else:
# No new data from the device; yield
# We cannot sleep(0) here because that would go 100% CPU if both queues are empty.
gevent.sleep(DEVICE_POLL_INTERVAL)
except KeyboardInterrupt:
self.running = False
hidraw.close()
gevent.kill(crypto, KeyboardInterrupt)
gevent.kill(console_updater, KeyboardInterrupt)
def setup_crypto(self, sn):
"""
Performs decryption of packets received. Stores decrypted packets in a Queue for use.
"""
if is_old_model(sn):
self.old_model = True
print self.old_model
k = ['\0'] * 16
k[0] = sn[-1]
k[1] = '\0'
k[2] = sn[-2]
if self.is_research:
k[3] = 'H'
k[4] = sn[-1]
k[5] = '\0'
k[6] = sn[-2]
k[7] = 'T'
k[8] = sn[-3]
k[9] = '\x10'
k[10] = sn[-4]
k[11] = 'B'
else:
k[3] = 'T'
k[4] = sn[-3]
k[5] = '\x10'
k[6] = sn[-4]
k[7] = 'B'
k[8] = sn[-1]
k[9] = '\0'
k[10] = sn[-2]
k[11] = 'H'
k[12] = sn[-3]
k[13] = '\0'
k[14] = sn[-4]
k[15] = 'P'
key = ''.join(k)
iv = Random.new().read(AES.block_size)
cipher = AES.new(key, AES.MODE_ECB, iv)
print ("Encryption...")
for i in k:
print "0x%.02x " % (ord(i))
while self.running:
print ("Check encryption queue...")
while not tasks.empty():
task = tasks.get()
try:
print ("Adding packets to queue...:"+str(self.packets.qsize()))
data = cipher.decrypt(task[:16]) + cipher.decrypt(task[16:])
print ("Data received:"+str(len(data)))
self.packets.put_nowait(EmotivPacket(data, self.sensors, self.old_model))
self.packets_processed += 1
print ("Packets added..")
except Exception as ex:
print ("Exception raied..:"+str(ex))
pass
gevent.sleep(0.01)
gevent.sleep(0.01)
def dequeue(self):
"""
Returns an EmotivPacket popped off the Queue.
"""
try:
print ("Returning something...:"+str(self.packets.empty()))
return self.packets.get()
except Exception, e:
print e
class Emotiv(object):
"""
Receives, decrypts and stores packets received from Emotiv Headsets.
"""
def __init__(self,
display_output=True,
serial_number="",
is_research=False):
"""
Sets up initial values.
"""
self.running = True
self.packets = Queue()
self.packets_received = 0
self.packets_processed = 0
self.battery = 0
self.display_output = display_output
self.is_research = is_research
self.sensors = {
'F3': {
'value': 0,
'quality': 0
},
'FC6': {
'value': 0,
'quality': 0
},
'P7': {
'value': 0,
'quality': 0
},
'T8': {
'value': 0,
'quality': 0
},
'F7': {
'value': 0,
'quality': 0
},
'F8': {
'value': 0,
'quality': 0
},
'T7': {
'value': 0,
'quality': 0
},
'P8': {
'value': 0,
'quality': 0
},
'AF4': {
'value': 0,
'quality': 0
},
'F4': {
'value': 0,
'quality': 0
},
'AF3': {
'value': 0,
'quality': 0
},
'O2': {
'value': 0,
'quality': 0
},
'O1': {
'value': 0,
'quality': 0
},
'FC5': {
'value': 0,
'quality': 0
},
'X': {
'value': 0,
'quality': 0
},
'Y': {
'value': 0,
'quality': 0
},
'Unknown': {
'value': 0,
'quality': 0
}
}
self.serial_number = serial_number # You will need to set this manually for OS X.
self.old_model = False
def setup(self):
"""
Runs setup function depending on platform.
"""
print(system_platform + " detected.")
if system_platform == "Windows":
self.setup_windows()
elif system_platform == "Linux":
self.setup_posix()
elif system_platform == "Darwin":
self.setup_darwin()
def setup_windows(self):
"""
Setup for headset on the Windows platform.
"""
devices = []
try:
for device in hid.find_all_hid_devices():
if device.vendor_id != 0x21A1 and device.vendor_id != 0xED02:
continue
if device.product_name == 'Brain Waves':
devices.append(device)
device.open()
self.serial_number = device.serial_number
device.set_raw_data_handler(self.handler)
elif device.product_name == 'EPOC BCI':
devices.append(device)
device.open()
self.serial_number = device.serial_number
device.set_raw_data_handler(self.handler)
elif device.product_name == '00000000000':
devices.append(device)
device.open()
self.serial_number = device.serial_number
device.set_raw_data_handler(self.handler)
elif device.product_name == 'Emotiv RAW DATA':
devices.append(device)
device.open()
self.serial_number = device.serial_number
device.set_raw_data_handler(self.handler)
crypto = gevent.spawn(self.setup_crypto, self.serial_number)
console_updater = gevent.spawn(self.update_console)
while self.running:
try:
gevent.sleep(0)
except KeyboardInterrupt:
self.running = False
finally:
for device in devices:
device.close()
gevent.kill(crypto, KeyboardInterrupt)
gevent.kill(console_updater, KeyboardInterrupt)
def handler(self, data):
"""
Receives packets from headset for Windows. Sends them to a Queue to be processed
by the crypto greenlet.
"""
assert data[0] == 0
tasks.put_nowait(''.join(map(chr, data[1:])))
self.packets_received += 1
return True
def setup_posix(self):
"""
Setup for headset on the Linux platform.
Receives packets from headset and sends them to a Queue to be processed
by the crypto greenlet.
"""
_os_decryption = False
if os.path.exists('/dev/eeg/raw'):
# The decryption is handled by the Linux epoc daemon. We don't need to handle it.
_os_decryption = True
hidraw = open("/dev/eeg/raw", 'rb')
else:
serial, hidraw_filename = get_linux_setup()
self.serial_number = serial
if os.path.exists("/dev/" + hidraw_filename):
hidraw = open("/dev/" + hidraw_filename, 'rb')
else:
hidraw = open("/dev/hidraw4", 'rb')
crypto = gevent.spawn(self.setup_crypto, self.serial_number)
console_updater = gevent.spawn(self.update_console)
while self.running:
try:
data = hidraw.read(32)
if data != "":
if _os_decryption:
self.packets.put_nowait(EmotivPacket(data))
else:
#Queue it!
self.packets_received += 1
tasks.put_nowait(data)
gevent.sleep(0)
else:
# No new data from the device; yield
# We cannot sleep(0) here because that would go 100% CPU if both queues are empty
gevent.sleep(DEVICE_POLL_INTERVAL)
except KeyboardInterrupt:
self.running = False
hidraw.close()
if not _os_decryption:
gevent.kill(crypto, KeyboardInterrupt)
gevent.kill(console_updater, KeyboardInterrupt)
def setup_darwin(self):
"""
Setup for headset on the OS X platform.
Receives packets from headset and sends them to a Queue to be processed
by the crypto greenlet.
"""
_os_decryption = False
# Change these values to the hex equivalent from the output of hid_enumerate. If they are incorrect.
# Current values = VendorID: 8609 ProductID: 1
hidraw = hid.device(0x21a1, 0x0001)
if not hidraw:
hidraw = hid.device(0x21a1, 0x1234)
if not hidraw:
hidraw = hid.device(0xed02, 0x1234)
if not hidraw:
print(
"Device not found. Uncomment the code in setup_darwin and modify hid.device(vendor_id, product_id)"
)
raise ValueError
if self.serial_number == "":
print(
"Serial number needs to be specified manually in __init__().")
raise ValueError
crypto = gevent.spawn(self.setup_crypto, self.serial_number)
console_updater = gevent.spawn(self.update_console)
zero = 0
while self.running:
try:
# Doesn't seem to matter how big we make the buffer 32 returned every time, 33 for other platforms
data = hidraw.read(34)
if len(data) == 32:
# Most of the time the 0 is truncated? That's ok we'll add it...
data = [zero] + data
if data != "":
if _os_decryption:
self.packets.put_nowait(EmotivPacket(data))
else:
#Queue it!
tasks.put_nowait(''.join(map(chr, data[1:])))
self.packets_received += 1
gevent.sleep(0)
else:
# No new data from the device; yield
# We cannot sleep(0) here because that would go 100% CPU if both queues are empty.
gevent.sleep(DEVICE_POLL_INTERVAL)
except KeyboardInterrupt:
self.running = False
hidraw.close()
gevent.kill(crypto, KeyboardInterrupt)
gevent.kill(console_updater, KeyboardInterrupt)
def setup_crypto(self, sn):
"""
Performs decryption of packets received. Stores decrypted packets in a Queue for use.
"""
if is_old_model(sn):
self.old_model = True
print(self.old_model)
k = ['\0'] * 16
k[0] = sn[-1]
k[1] = '\0'
k[2] = sn[-2]
if self.is_research:
k[3] = 'H'
k[4] = sn[-1]
k[5] = '\0'
k[6] = sn[-2]
k[7] = 'T'
k[8] = sn[-3]
k[9] = '\x10'
k[10] = sn[-4]
k[11] = 'B'
else:
k[3] = 'T'
k[4] = sn[-3]
k[5] = '\x10'
k[6] = sn[-4]
k[7] = 'B'
k[8] = sn[-1]
k[9] = '\0'
k[10] = sn[-2]
k[11] = 'H'
k[12] = sn[-3]
k[13] = '\0'
k[14] = sn[-4]
k[15] = 'P'
key = ''.join(k)
iv = Random.new().read(AES.block_size)
cipher = AES.new(key, AES.MODE_ECB, iv)
for i in k:
print("0x%.02x " % (ord(i)))
while self.running:
while not tasks.empty():
task = tasks.get()
data = cipher.decrypt(task[:16]) + cipher.decrypt(task[16:])
self.packets.put_nowait(
EmotivPacket(data, self.sensors, self.old_model))
self.packets_processed += 1
gevent.sleep(0)
gevent.sleep(0)
def dequeue(self):
"""
Returns an EmotivPacket popped off the Queue.
"""
try:
return self.packets.get()
except Exception as e:
print(e)
def close(self):
"""
Shuts down the running greenlets.
"""
self.running = False
def update_console(self):
"""
Greenlet that outputs sensor, gyro and battery values once per second to the console.
"""
if self.display_output:
while self.running:
if system_platform == "Windows":
os.system('cls')
else:
os.system('clear')
print("Packets Received: %s Packets Processed: %s" %
(self.packets_received, self.packets_processed))
print('\n'.join("%s Reading: %s Quality: %s" %
(k[1], self.sensors[k[1]]['value'],
self.sensors[k[1]]['quality'])
for k in enumerate(self.sensors)))
print("Battery: %i" % g_battery)
gevent.sleep(.001)
t = Tags(x)
g = pool.spawn(t.work)
threads.append(g)
gevent.joinall(threads)
if not os.path.isfile(tag_result_file):
l = []
for x in os.listdir(tag_result_path):
url = tag_result_path + x
r = [x.replace('\n', '') for x in open(url, 'r')]
l += r
l = set(l)
fw = io.open(tag_result_file, 'w', encoding='utf-8')
for x in l:
fw.write(unicode(x + '\n'))
fw.close()
# 2. 访问并保存 subject
q = Queue()
for x in open(tag_result_file, 'r'):
q.put_nowait(str(x).replace('\n', ''))
pool = Pool(args.thread)
threads = []
while not q.empty():
url = q.get()
s = Subject(url)
t = pool.spawn(s.work)
threads.append(t)
gevent.joinall(threads)
class Connection(object):
in_buffer_size = 4096
out_buffer_size = 4096
cql_version = None
keyspace = None
compression = True
compressor = None
decompressor = None
ssl_options = None
last_error = None
in_flight = 0
is_defunct = False
is_closed = False
lock = None
def __init__(
self,
host="127.0.0.1",
port=9042,
credentials=None,
ssl_options=None,
sockopts=None,
compression=True,
cql_version=None,
):
self.host = host
self.port = port
self.credentials = credentials
self.ssl_options = ssl_options
self.sockopts = sockopts
self.compression = compression
self.cql_version = cql_version
self._id_queue = Queue(MAX_STREAM_PER_CONNECTION)
for i in range(MAX_STREAM_PER_CONNECTION):
self._id_queue.put_nowait(i)
self.lock = RLock()
def close(self):
raise NotImplementedError()
def defunct(self, exc):
with self.lock:
if self.is_defunct or self.is_closed:
return
self.is_defunct = True
trace = traceback.format_exc(exc)
if trace != "None":
log.debug("Defuncting connection (%s) to %s: %s\n%s", id(self), self.host, exc, traceback.format_exc(exc))
else:
log.debug("Defuncting connection (%s) to %s: %s", id(self), self.host, exc)
self.last_error = exc
self.close()
self.error_all_callbacks(exc)
self.connected_event.set()
return exc
def error_all_callbacks(self, exc):
with self.lock:
callbacks = self._callbacks
self._callbacks = {}
new_exc = ConnectionShutdown(str(exc))
for cb in callbacks.values():
try:
cb(new_exc)
except Exception:
log.warn(
"Ignoring unhandled exception while erroring callbacks for a " "failed connection (%s) to host %s:",
id(self),
self.host,
exc_info=True,
)
def handle_pushed(self, response):
log.debug("Message pushed from server: %r", response)
for cb in self._push_watchers.get(response.event_type, []):
try:
cb(response.event_args)
except Exception:
log.exception("Pushed event handler errored, ignoring:")
def send_msg(self, msg, cb, wait_for_id=False):
if self.is_defunct:
raise ConnectionShutdown("Connection to %s is defunct" % self.host)
elif self.is_closed:
raise ConnectionShutdown("Connection to %s is closed" % self.host)
if not wait_for_id:
try:
request_id = self._id_queue.get_nowait()
except Queue.Empty:
raise ConnectionBusy("Connection to %s is at the max number of requests" % self.host)
else:
request_id = self._id_queue.get()
self._callbacks[request_id] = cb
self.push(msg.to_string(request_id, compression=self.compressor))
return request_id
def wait_for_response(self, msg, timeout=None):
return self.wait_for_responses(msg, timeout=timeout)[0]
def wait_for_responses(self, *msgs, **kwargs):
timeout = kwargs.get("timeout")
waiter = ResponseWaiter(self, len(msgs))
# busy wait for sufficient space on the connection
messages_sent = 0
while True:
needed = len(msgs) - messages_sent
with self.lock:
available = min(needed, MAX_STREAM_PER_CONNECTION - self.in_flight)
self.in_flight += available
for i in range(messages_sent, messages_sent + available):
self.send_msg(msgs[i], partial(waiter.got_response, index=i), wait_for_id=True)
messages_sent += available
if messages_sent == len(msgs):
break
else:
if timeout is not None:
timeout -= 0.01
if timeout <= 0.0:
raise OperationTimedOut()
time.sleep(0.01)
try:
return waiter.deliver(timeout)
except OperationTimedOut:
raise
except Exception, exc:
self.defunct(exc)
raise
class Socket(object):
"""
Virtual Socket implementation, checks heartbeats, writes to local queues
for message passing, holds the Namespace objects, dispatches de packets
to the underlying namespaces.
This is the abstraction on top of the different transports. It's like
if you used a WebSocket only...
"""
STATE_CONNECTING = "CONNECTING"
STATE_CONNECTED = "CONNECTED"
STATE_DISCONNECTING = "DISCONNECTING"
STATE_DISCONNECTED = "DISCONNECTED"
GLOBAL_NS = ''
"""Use this to be explicit when specifying a Global Namespace (an endpoint
with no name, not '/chat' or anything."""
def __init__(self, server, error_handler=None):
self.server = weakref.proxy(server)
self.sessid = str(random.random())[2:]
self.session = {} # the session dict, for general developer usage
self.client_queue = Queue() # queue for messages to client
self.server_queue = Queue() # queue for messages to server
self.hits = 0
self.heartbeats = 0
self.timeout = Event()
self.wsgi_app_greenlet = None
self.state = "NEW"
self.connection_confirmed = False
self.ack_callbacks = {}
self.ack_counter = 0
self.request = None
self.environ = None
self.namespaces = {}
self.active_ns = {} # Namespace sessions that were instantiated
self.jobs = []
self.error_handler = default_error_handler
if error_handler is not None:
self.error_handler = error_handler
def _set_namespaces(self, namespaces):
"""This is a mapping (dict) of the different '/namespaces' to their
BaseNamespace object derivative.
This is called by socketio_manage()."""
self.namespaces = namespaces
def _set_request(self, request):
"""Saves the request object for future use by the different Namespaces.
This is called by socketio_manage().
"""
self.request = request
def _set_environ(self, environ):
"""Save the WSGI environ, for future use.
This is called by socketio_manage().
"""
self.environ = environ
def _set_error_handler(self, error_handler):
"""Changes the default error_handler function to the one specified
This is called by socketio_manage().
"""
self.error_handler = error_handler
def _get_next_msgid(self):
"""This retrieves the next value for the 'id' field when sending
an 'event' or 'message' or 'json' that asks the remote client
to 'ack' back, so that we trigger the local callback.
"""
self.ack_counter += 1
return self.ack_counter
def _save_ack_callback(self, msgid, callback):
"""Keep a reference of the callback on this socket."""
if msgid in self.ack_callbacks:
return False
self.ack_callbacks[msgid] = callback
def _pop_ack_callback(self, msgid):
"""Fetch the callback for a given msgid, if it exists, otherwise,
return None"""
if msgid not in self.ack_callbacks:
return None
return self.ack_callbacks.pop(msgid)
def __str__(self):
result = ['sessid=%r' % self.sessid]
if self.state == self.STATE_CONNECTED:
result.append('connected')
if self.client_queue.qsize():
result.append('client_queue[%s]' % self.client_queue.qsize())
if self.server_queue.qsize():
result.append('server_queue[%s]' % self.server_queue.qsize())
if self.hits:
result.append('hits=%s' % self.hits)
if self.heartbeats:
result.append('heartbeats=%s' % self.heartbeats)
return ' '.join(result)
def __getitem__(self, key):
"""This will get the nested Namespace using its '/chat' reference.
Using this, you can go from one Namespace to the other (to emit, add
ACLs, etc..) with:
adminnamespace.socket['/chat'].add_acl_method('kick-ban')
"""
return self.active_ns[key]
def __hasitem__(self, key):
"""Verifies if the namespace is active (was initialized)"""
return key in self.active_ns
@property
def connected(self):
"""Returns whether the state is CONNECTED or not."""
return self.state == self.STATE_CONNECTED
def incr_hits(self):
self.hits += 1
if self.hits == 1:
self.state = self.STATE_CONNECTED
def heartbeat(self):
"""This makes the heart beat for another X seconds. Call this when
you get a heartbeat packet in.
This clear the heartbeat disconnect timeout (resets for X seconds).
"""
self.timeout.set()
def kill(self):
"""This function must/will be called when a socket is to be completely
shut down, closed by connection timeout, connection error or explicit
disconnection from the client.
It will call all of the Namespace's
:meth:`~socketio.namespace.BaseNamespace.disconnect` methods
so that you can shut-down things properly.
"""
# Clear out the callbacks
self.ack_callbacks = {}
if self.connected:
self.state = self.STATE_DISCONNECTING
self.server_queue.put_nowait(None)
self.client_queue.put_nowait(None)
self.disconnect()
# gevent.kill(self.wsgi_app_greenlet)
else:
pass # Fail silently
def put_server_msg(self, msg):
"""Writes to the server's pipe, to end up in in the Namespaces"""
self.heartbeat()
self.server_queue.put_nowait(msg)
def put_client_msg(self, msg):
"""Writes to the client's pipe, to end up in the browser"""
self.heartbeat()
self.client_queue.put_nowait(msg)
def get_client_msg(self, **kwargs):
"""Grab a message to send it to the browser"""
return self.client_queue.get(**kwargs)
def get_server_msg(self, **kwargs):
"""Grab a message, to process it by the server and dispatch calls
"""
return self.server_queue.get(**kwargs)
def get_multiple_client_msgs(self, **kwargs):
"""Get multiple messages, in case we're going through the various
XHR-polling methods, on which we can pack more than one message if the
rate is high, and encode the payload for the HTTP channel."""
client_queue = self.client_queue
msgs = [client_queue.get(**kwargs)]
while client_queue.qsize():
msgs.append(client_queue.get())
return msgs
def error(self, error_name, error_message, endpoint=None, msg_id=None,
quiet=False):
"""Send an error to the user, using the custom or default
ErrorHandler configured on the [TODO: Revise this] Socket/Handler
object.
:param error_name: is a simple string, for easy association on
the client side
:param error_message: is a human readable message, the user
will eventually see
:param endpoint: set this if you have a message specific to an
end point
:param msg_id: set this if your error is relative to a
specific message
:param quiet: way to make the error handler quiet. Specific to
the handler. The default handler will only log,
with quiet.
"""
handler = self.error_handler
return handler(self, error_name, error_message, endpoint, msg_id, quiet)
# User facing low-level function
def disconnect(self, silent=False):
"""Calling this method will call the
:meth:`~socketio.namespace.BaseNamespace.disconnect` method on
all the active Namespaces that were open, killing all their
jobs and sending 'disconnect' packets for each of them.
Normally, the Global namespace (endpoint = '') has special meaning,
as it represents the whole connection,
:param silent: when True, pass on the ``silent`` flag to the Namespace
:meth:`~socketio.namespace.BaseNamespace.disconnect` calls.
"""
for ns_name, ns in list(self.active_ns.iteritems()):
ns.disconnect(silent=silent)
def remove_namespace(self, namespace):
"""This removes a Namespace object from the socket.
This is usually called by
:meth:`~socketio.namespace.BaseNamespace.disconnect`.
"""
if namespace in self.active_ns:
del self.active_ns[namespace]
def send_packet(self, pkt):
"""Low-level interface to queue a packet on the wire (encoded as wire
protocol"""
self.put_client_msg(packet.encode(pkt))
def spawn(self, fn, *args, **kwargs):
"""Spawn a new Greenlet, attached to this Socket instance.
It will be monitored by the "watcher" method
"""
self.debug("Spawning sub-Socket Greenlet: %s" % fn.__name__)
job = gevent.spawn(fn, *args, **kwargs)
self.jobs.append(job)
return job
def _receiver_loop(self):
"""This is the loop that takes messages from the queue for the server
to consume, decodes them and dispatches them.
"""
while True:
rawdata = self.get_server_msg()
if not rawdata:
continue # or close the connection ?
try:
pkt = packet.decode(rawdata)
except (ValueError, KeyError, Exception), e:
self.error('invalid_packet', "There was a decoding error when dealing with packet with event: %s... (%s)" % (rawdata[:20], e))
continue
if pkt['type'] == 'heartbeat':
# This is already dealth with in put_server_msg() when
# any incoming raw data arrives.
continue
if pkt['type'] == 'disconnect' and pkt['endpoint'] == '':
# On global namespace, we kill everything.
self.kill()
continue
endpoint = pkt['endpoint']
if endpoint not in self.namespaces:
self.error("no_such_namespace", "The endpoint you tried to connect to doesn't exist: %s" % endpoint, endpoint=endpoint)
continue
elif endpoint in self.active_ns:
pkt_ns = self.active_ns[endpoint]
else:
new_ns_class = self.namespaces[endpoint]
pkt_ns = new_ns_class(self.environ, endpoint,
request=self.request)
pkt_ns.initialize() # use this instead of __init__,
# for less confusion
self.active_ns[endpoint] = pkt_ns
retval = pkt_ns.process_packet(pkt)
# Has the client requested an 'ack' with the reply parameters ?
if pkt.get('ack') == "data" and pkt.get('id'):
returning_ack = dict(type='ack', ackId=pkt['id'],
args=retval,
endpoint=pkt.get('endpoint', ''))
self.send_packet(returning_ack)
# Now, are we still connected ?
if not self.connected:
self.kill() # ?? what,s the best clean-up when its not a
# user-initiated disconnect
return
from gevent import monkey
monkey.patch_all() #让程序变成异步模式
from gevent.queue import Queue
from bs4 import BeautifulSoup
from openpyxl import Workbook
from openpyxl.writer.excel import ExcelWriter
from openpyxl import load_workbook
import requests, gevent, csv
work = Queue()
url_1 = 'http://www.boohee.com/food/group/{type}?page={page}'
url_2 = 'http://www.boohee.com/food/view_menu?page={page}'
for i in range(1, 3):
for j in range(1, 3):
real_url = url_1.format(type=i, page=j)
work.put_nowait(real_url)
for x in range(3):
real_url = url_2.format(page=x)
work.put_nowait(real_url)
def boohe_spider():
headers = {
'user-agent':
'Mozilla/5.0 (Windows NT 6.1; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/78.0.3904.87 Safari/537.36'
}
while not work.empty(): #当队列不为空，执行下面的代码
url = work.get_nowait() #使用get_nowait()方法将请求url从队列中取出来
response = requests.get(url=url, headers=headers)
print('第{}个食物的第{}页'.format(i, j))
class Worker(object):
def __init__(self, master_ip_port, worker_ip_port):
self.state = STATE_READY
self.master_ip_port = master_ip_port
self.worker_ip_port = worker_ip_port
self.map_object = None
self.reduce_object = None
self.unit = None
self.data_dir = None
self.reducer_map_queue = Queue()
self.master_c = None
gevent.spawn(self.controller)
def controller(self):
while True:
print('[Worker: %s]: %s' % (self.worker_ip_port, self.state))
gevent.sleep(1)
def ping(self):
print('[Worker] Ping from Master')
return True
# Methods for mapper
def do_map(self, split_info, task_file, mapper_id, num_mapper, num_reducer, task_name, file_info, split_size):
gevent.spawn(self.do_map_async, split_info,
task_file, mapper_id, num_mapper, num_reducer, task_name, file_info, split_size)
def do_map_async(self, split_info,
task_file, mapper_id, num_mapper, num_reducer, task_name, file_info, split_size):
print "Get task %s from master, current mapper id is %d" % (task_name, mapper_id)
self.read_map_task(task_file, mapper_id, num_mapper, num_reducer, task_name, split_size)
print "Task %s M %d: read user map function successfully" % (task_name, mapper_id)
data = self.read_input(split_info, mapper_id, num_mapper, file_info, split_size)
print "Task %s M %d: read input successfully" % (task_name, mapper_id)
# Ensure worker can get master ping
# gevent.sleep(0)
self.map_object.map(task_file[0], data)
print "Task %s M %d: do map successfully" % (task_name, mapper_id)
if test_failure:
gevent.sleep(3)
self.map_object.combine()
print "Task %s M %d: combine intermediate result successfully" % (task_name, mapper_id)
# Ensure worker can get master ping
# gevent.sleep(0)
self.map_object.write_to_file()
print "Task %s M %d: write result to file successfully" % (task_name, mapper_id)
temp_script = task_file[0] + "m" + str(mapper_id)
os.remove(temp_script)
print "Task %s M %d: delete temp_script %s" % (task_name, mapper_id, temp_script)
print "Task %s M %d: Finished" % (task_name, mapper_id)
self.master_c.mapper_finish(True, task_name, mapper_id, self.worker_ip_port)
# read the data from split and also keep unit ( i.e. get next line from next split)
# split_info : single file [(file_name, start, end)]
# or multiple files [(file_name0, start, end), (file_name1, start, end)]
# mapper_id : id of this mapper
# num_mapper : how many reducers is in this task
# file_info : all files info in this task
# [(file0_path, file0_size), (file1_path, file1_size)]
# split_size : the length of split data
def read_input(self, split_info, mapper_id, num_mapper, file_info, split_size):
data = ""
filename = ""
start = 0
read_size = 0
if debug:
print "From worker mapper %d split_info" % mapper_id, split_info
# read data from the split_info for this mapper
for file in split_info:
filename = file[0]
start = file[1]
read_size = file[2] - file[1]
data += self.read_data_from_file(filename, start, read_size)
last_file_path = filename
start = start + read_size
# get the last filename of this mapper in file_info
used_file = 0
for file in file_info:
if file[0] == last_file_path:
break
used_file += 1
if used_file > len(file_info):
raise Exception("can't find the last file in split")
# remove the data fetch by the previous mapper and get more data to keep unit
# If unit is int, then data == the multiple of unit
if type(self.unit) == int:
# Remove the unit data fetched by the previous mapper if this mapper is not the first
remove_size = 0
for i in range(mapper_id):
current_size = split_size - remove_size
if (current_size % self.unit) == 0:
remaining_size = 0
else:
remaining_size = self.unit - (current_size % self.unit)
remove_size = remaining_size
data = data[remove_size:]
current_size = split_size - remove_size
if (current_size % self.unit) == 0:
remaining_size = 0
else:
remaining_size = self.unit - (current_size % self.unit)
if debug:
print "remove_size %d" % remove_size
print "start %d remaining_size %d current_size %d" % (start, remaining_size, current_size)
print "data ", data
# Get more data ( if not the last mapper )
if mapper_id != num_mapper - 1:
while remaining_size > 0:
current_file_name = file_info[used_file][0]
current_file_size = file_info[used_file][1]
# Required remaining_size <= file remaining_size
if remaining_size <= (current_file_size - start):
data += self.read_data_from_file(current_file_name, start, remaining_size)
if remaining_size == current_file_size - start:
start = 0
used_file += 1
else:
start = start + remaining_size
remaining_size = 0
# Required remaining_size > file remaining_size
else:
if used_file < len(file_info) - 1:
data += self.read_data_from_file(current_file_name, start, current_file_size - start)
remaining_size -= current_file_size - start
start = 0
used_file += 1
# This is the last file, then finish split
else:
data += self.read_data_from_file(current_file_name, start, current_file_size - start)
remaining_size = 0
# If unit is str, then get more data until we see the unit
elif type(self.unit) == str:
# Remove the first split if mapper_id is not 0
if mapper_id != 0:
if len(data.split(self.unit)) > 1:
data = data.split(self.unit, 1)[1]
else:
data = ""
# Get more split if the mapper is not the last mapper
if mapper_id != num_mapper - 1:
data += self.read_data_from_file(file_info[used_file][0], start, file_info[used_file][1] - start) \
.split(self.unit, 1)[0]
if debug:
print "From worker mapper %d input data" % mapper_id, data
return data
# read data from file
def read_data_from_file(self, filename, start, read_size):
f = open(filename)
f.seek(start)
data = f.read(read_size)
try:
f.close()
except:
print "Error: can't close the original data file"
return data
# read map_task from task_file,
# task_file : (task_filename, code)
# mapper_id : this mapper #
# num_mapper : how many mappers is in this task
# num_reducer : how many reducers is in this task
# task_name : task name( it's the time that the client submit the job)
# split_size : how many bytes is in one split
def read_map_task(self, task_file, mapper_id, num_mapper, num_reducer, task_name, split_size):
f = open(task_file[0] + 'm' + str(mapper_id), 'w')
f.write(task_file[1])
try:
f.close()
except:
print "Error: can't close the map script"
module = importlib.import_module(task_file[0][:-3])
self.unit = module.unit
self.map_object = module.Map(mapper_id, num_mapper, num_reducer, task_name, split_size)
# reducer_id want to fetch intermediate file(mapper_id) from this mapper
# but this current mapper_id can be different from mapper_id
def fetch_intermediate_file(self, task_name, mapper_id, reducer_id):
table = pickle.load(open(task_name + "_m" + str(mapper_id) + "_r" + str(reducer_id), 'rb'))
return table
def remove_intermediate_file(self, task_name, mapper_id, num_reducer):
for reducer_id in range(num_reducer):
intermediate_filename = task_name + '_m' + str(mapper_id) + "_r" + str(reducer_id)
os.remove(intermediate_filename)
# master want to fetch output file from this reducer
def fetch_result_file(self, output_file, reducer_id):
output_filename = output_file + '_' + str(reducer_id)
f = open(output_filename, 'rb')
result = f.read()
try:
f.close()
except:
print "Error: can't close intermediate file"
os.remove(output_filename)
return result
# Methods for reducer
def do_reduce(self, task_file, num_mapper, reducer_id, output_file, task_name):
gevent.spawn(self.do_reduce_async, task_file, num_mapper, reducer_id, output_file, task_name)
def do_reduce_async(self,task_file, num_mapper, reducer_id, output_file, task_name):
print "Get task %s from master, current reducer id is %d" % (task_name, reducer_id)
self.read_reduce_task(task_file, reducer_id, output_file)
print "Task %s R %d: read user reduce function successfully" % (task_name, reducer_id)
entire_data = {}
while len(entire_data) != num_mapper:
# Ensure reducer get master ping()
# gevent.sleep(0)
try:
mapper_id, mapper_ip_port = self.reducer_map_queue.get()
except:
print "**** Error: can't get enough mapper intermediate file"
return
if mapper_ip_port == self.worker_ip_port:
try:
data = pickle.load(open(task_name + "_m" + str(mapper_id) +
'_r' + str(reducer_id), 'rb'))
except:
print "**** Error: can't get intermediate file from local mapper %d %s" \
% (mapper_id, mapper_ip_port)
else:
try:
c = zerorpc.Client()
c.connect("tcp://" + mapper_ip_port)
data = c.fetch_intermediate_file(task_name, mapper_id, reducer_id)
c.close()
except:
print "**** Error: can't get intermediate file from mapper %d %s" \
% (mapper_id, mapper_ip_port)
entire_data[mapper_id] = data
print "Task %s R %d: get intermediate data from mapper %d: %s" \
% (task_name, reducer_id, mapper_id, mapper_ip_port)
self.reduce_object.reduce(task_name, entire_data)
print "Task %s R %d: do reduce successfully" % (task_name, reducer_id)
self.reduce_object.write_to_file()
print "Task %s R %d: write result to file successfully" % (task_name, reducer_id)
temp_script = task_file[0] + "r" + str(reducer_id)
os.remove(temp_script)
print "Task %s R %d: delete temp_script %s" % (task_name, reducer_id, temp_script)
print "Task %s R %d: Finished" % (task_name, reducer_id)
self.master_c.reducer_finish(True, task_name, reducer_id, self.worker_ip_port)
# reducer is notified by master that the mapper has finished its job
def notify_mapper_finish(self, mapper_id, mapper_ip_port):
self.reducer_map_queue.put_nowait((mapper_id, mapper_ip_port))
# Read reduce_task from task_file
# task_file = (task_filename, code)
def read_reduce_task(self, task_file, reduce_id, output_file):
f = open(task_file[0] + 'r' + str(reduce_id), 'w')
f.write(task_file[1])
try:
f.close()
except:
print "Error: can't close the reduce script"
module = importlib.import_module(task_file[0][:-3])
self.reduce_object = module.Reduce(reduce_id, output_file)
from gevent import monkey
monkey.patch_all()
import gevent, time, requests, csv
from gevent.queue import Queue
from bs4 import BeautifulSoup
work = Queue()
cvs_file = open("Top100.csv", 'w', newline='', encoding='utf-8')
for i in range(10):
if i == 0:
work.put_nowait('http://www.mtime.com/top/tv/top100/index.html')
else:
url = 'http://www.mtime.com/top/tv/top100/index-%s.html' % (i + 1)
work.put_nowait(url)
def crawler():
while not work.empty():
url = work.get_nowait()
req = requests.get(url)
soup = BeautifulSoup(req.text, 'html.parser')
h2_list = soup.find_all("h2", class_="px14 pb6")
writer = csv.writer(cvs_file)
for h2 in h2_list:
name = h2.find('a').text
writer.writerow([name])
w_file = open('D:\\time_movie.csv', 'w', newline='', encoding='utf-8')
write = csv.writer(w_file)
row_header = ['剧名', '简介']
write.writerow(row_header)
work = Queue()
for x in range(10):
if x == 0:
x = ''
url = 'http://www.mtime.com/top/tv/top100'
else:
url = 'http://www.mtime.com/top/tv/top100/index-' + str(x +
1) + '.html'
work.put_nowait(url)
def crawler():
while not work.empty():
url = work.get_nowait()
res = requests.get(url)
bs = bs4.BeautifulSoup(res.text, 'html.parser')
bs = bs.find('ul', id="asyncRatingRegion")
for key in bs.find_all('li'):
title = key.find('h2', class_="px14 pb6").text
if key.find('p', class_="mt3") != None:
info = key.find('p', class_="mt3").text
write.writerow([title, info])
class ScraperGeventQueue(object):
'''
A gevent queue for carrefourScraper.
Parameters
----------
scraperClass: class
carrefour page scraper for CPP
args: list
a multi-list for CPP
gevent_num=100: int
maximum running gevent number.
'''
def __init__(self, scraperClass, args, gevent_num=80):
self.scraperClass = scraperClass
self.args = self._args(args)
self.gevent_num = gevent_num
self.tasks = Queue() # create a gevent queue
self.failure_list = []
def _args(self, args):
'''
Parse input args.
Parameters
----------
category: list or tuple
a group of input categories
areas: list or tuple
a group pf input area information
Returns
-------
new_args: list
a group of new input args
'''
new_args = []
for arg in args:
category, selectedCategoryId = splitPara(arg)
new_args.append((category, selectedCategoryId, 1))
return new_args
def _run(self, task):
'''
Define the run function.
Parameters
----------
task: list or tuple
a group of parameters for self.scraperClass
'''
# split parameters for the page scraper class
category_name, selectedCategoryId, page_num = task
# run the page scraper class
scraper = self.scraperClass(category_name, selectedCategoryId,
page_num)
scraper.json = scraper.getJSON()
# if connect error, add the parameters into queue once again
if (not scraper.json) and (task not in self.failure_list):
self.failure_list.append(task)
self.tasks.put_nowait(task)
data_list = scraper.parseJSON(scraper.json)
#indicator = scraper.writeMongoDB(data_list)
indicator = scraper.writeCSV(data_list)
# produce new parameters and add them to the gevent queue
if (page_num == 1) and indicator:
total_page = scraper.getTotalPageNumber()
if total_page > 1:
[
self.tasks.put_nowait(
(category_name, selectedCategoryId, i))
for i in range(2, total_page + 1)
]
def worker(self):
'A gevent worker.'
while not self.tasks.empty():
task = self.tasks.get()
self._run(task)
def manager(self):
'A gevent manager, creating the initial gevents'
for arg in self.args:
self.tasks.put_nowait(arg)
def start(self):
'Run the gevent queue.'
gevent.spawn(self.manager).join()
tasks = [gevent.spawn(self.worker) for i in range(self.gevent_num)]
gevent.joinall(tasks)
class Session(object):
"""
Base class for SockJS sessions. Provides a transport independent way to
queue messages from/to the client.
``state``: Session state
``manager``: Session manager that hold this session
``acquired``: Acquired state, indicates that transport is using session
``timeout``: Session timeout
"""
__slots__ = ('id', 'acquired', 'state', 'interrupted', 'exception',
'manager', 'handler', 'expired', 'timeout', 'expires',
'_hits', '_heartbeats', '_heartbeat_transport', '_debug',
'_waiter', '_queue', 'registry', '__weakref__')
def __init__(self, id, handler, timeout=timedelta(seconds=5), debug=True):
# Protected
self.manager = None
self.acquired = False
self.interrupted = False
self.exception = None
# Private
self._hits = 0
self._heartbeats = 0
self._heartbeat_transport = False
self._debug = debug
self._waiter = None
self._queue = Queue()
# Public
self.id = id
self.handler = handler
self.state = STATE_NEW
self.expired = False
self.timeout = timeout
self.expires = datetime.now() + timeout
def __str__(self):
result = ['id=%r' % (self.id, )]
if self.state == STATE_OPEN:
result.append('connected')
elif self.state == STATE_CLOSED:
result.append('closed')
else:
result.append('disconnected')
if self.acquired:
result.append('acquired')
if len(self._queue):
result.append('queue[%s]' % len(self._queue))
if self._hits:
result.append('hits=%s' % self._hits)
if self._heartbeats:
result.append('heartbeats=%s' % self._heartbeats)
return ' '.join(result)
def _tick(self, timeout=None):
""" Bump the TTL of the session. """
log.info("Update expires time {} for session".format(
self.expires, self.id))
self.expired = False
if timeout is None:
self.expires = datetime.now() + self.timeout
else:
self.expires = datetime.now() + timeout
def _acquire(self, manager=None, heartbeat=True):
self.acquired = True
self.manager = manager
self._heartbeat_transport = heartbeat
self._tick()
self._hits += 1
if self.state == STATE_NEW:
log.info('[sockjs_flask] Open session: %s', self.id)
self.state = STATE_OPEN
self.add_message(FRAME_OPEN, FRAME_OPEN)
try:
self.handler(OpenMessage, self)
except Exception as exc:
self.state = STATE_CLOSING
self.exception = exc
self.interrupted = True
self.add_message(FRAME_CLOSE, *CONN_CLOSED)
log.exception('Exception in open session handling.')
def _release(self):
log.info('[sockjs_flask] Session: %s release on %s', self.id,
datetime.now())
self.acquired = False
self.manager = None
self._heartbeat_transport = False
def _heartbeat(self):
self.expired = False
self._tick()
self._heartbeats += 1
self.add_message(FRAME_HEARTBEAT, FRAME_HEARTBEAT)
def _wait(self, pack=True):
"""
Get packet from queue and return frame pack with data
:param pack: Type return package
:return: tuple
"""
#log.info("Waiter {} session and queue {}".format(self._waiter, self._queue))
if self._queue:
frame, payload = self._queue.get()
if pack:
if frame == FRAME_CLOSE:
return FRAME_CLOSE, close_frame(*payload)
elif frame == FRAME_MESSAGE:
return FRAME_MESSAGE, messages_frame(payload)
return frame, payload
else:
raise SessionIsClosed()
def _remote_close(self, exc=None):
""" Close session from remote. """
if self.state in (STATE_CLOSING, STATE_CLOSED):
return
log.info('close session: %s', self.id)
self.state = STATE_CLOSING
if exc is not None:
self.exception = exc
self.interrupted = True
try:
self.handler(SockjsMessage(MSG_CLOSE, exc), self)
except Exception as e:
log.exception('Exception in close handler.')
def _remote_closed(self):
if self.state == STATE_CLOSED:
return
log.info('session closed: %s', self.id)
self.state = STATE_CLOSED
self.expire()
try:
self.handler(ClosedMessage, self)
except Exception as e:
log.exception('Exception in closed handler.')
waiter = self._waiter
if waiter is not None:
self._waiter = None
if not waiter.cancelled():
waiter.set_result(True)
def _remote_message(self, msg):
log.debug('incoming message: %s, %s', self.id, msg[:200])
self._tick()
try:
self.handler(SockjsMessage(MSG_MESSAGE, msg), self)
except:
log.exception('Exception in message handler.')
def add_message(self, frame, data):
log.info('session closed: %s', self.id)
self._queue.put_nowait((frame, data))
waiter = self._waiter
if waiter is not None:
self._waiter = None
if not waiter.cancelled():
waiter.set_result(True)
def expire(self):
self.expired = True
def send(self, msg):
""" Send message """
assert isinstance(msg, str), 'String is required'
if self._debug:
log.info('Оutgoing message from send: %s, %s', self.id,
str(msg)[:200])
if self.state != STATE_OPEN:
return
self._tick()
self.add_message(FRAME_MESSAGE, msg)
def send_frame(self, frm):
""" Send message frame to client """
if self._debug:
log.info('Оutgoing message from send_frame: %s, %s', self.id,
frm[:200])
if self.state != STATE_OPEN:
return
self._tick()
self.add_message(FRAME_MESSAGE_BLOB, frm)
def close(self, code=3000, reason='Go away!'):
""" Close session """
if self.state not in (STATE_CLOSING, STATE_CLOSED):
self.state = STATE_CLOSING
self.add_message(FRAME_CLOSE, (code, reason))
if self._debug:
log.debug('[sockjs_flask] Session closed: %s', self.id)
class Emotiv(object):
"""
Receives, decrypts and stores packets received from Emotiv Headsets.
"""
def __init__(self, display_output=True, serial_number="", is_research=False):
"""
Sets up initial values.
"""
self.running = True
self.packets = Queue()
self.packets_received = 0
self.packets_processed = 0
self.battery = 0
self.display_output = display_output
self.is_research = is_research
self.sensors = {
"F3": {"value": 0, "quality": 0},
"FC6": {"value": 0, "quality": 0},
"P7": {"value": 0, "quality": 0},
"T8": {"value": 0, "quality": 0},
"F7": {"value": 0, "quality": 0},
"F8": {"value": 0, "quality": 0},
"T7": {"value": 0, "quality": 0},
"P8": {"value": 0, "quality": 0},
"AF4": {"value": 0, "quality": 0},
"F4": {"value": 0, "quality": 0},
"AF3": {"value": 0, "quality": 0},
"O2": {"value": 0, "quality": 0},
"O1": {"value": 0, "quality": 0},
"FC5": {"value": 0, "quality": 0},
"X": {"value": 0, "quality": 0},
"Y": {"value": 0, "quality": 0},
"Unknown": {"value": 0, "quality": 0},
}
self.serial_number = serial_number # You will need to set this manually for OS X.
self.old_model = False
def setup(self):
"""
Runs setup function depending on platform.
"""
print system_platform + " detected."
if system_platform == "Windows":
self.setup_windows()
elif system_platform == "Linux":
self.setup_posix()
elif system_platform == "Darwin":
self.setup_darwin()
def setup_windows(self):
"""
Setup for headset on the Windows platform.
"""
devices = []
try:
for device in hid.find_all_hid_devices():
if device.vendor_id != 0x21A1 and device.vendor_id != 0xED02:
continue
if device.product_name == "Brain Waves":
devices.append(device)
device.open()
self.serial_number = device.serial_number
device.set_raw_data_handler(self.handler)
elif device.product_name == "EPOC BCI":
devices.append(device)
device.open()
self.serial_number = device.serial_number
device.set_raw_data_handler(self.handler)
elif device.product_name == "00000000000":
devices.append(device)
device.open()
self.serial_number = device.serial_number
device.set_raw_data_handler(self.handler)
elif device.product_name == "Emotiv RAW DATA":
devices.append(device)
device.open()
self.serial_number = device.serial_number
device.set_raw_data_handler(self.handler)
crypto = gevent.spawn(self.setup_crypto, self.serial_number)
console_updater = gevent.spawn(self.update_console)
while self.running:
try:
gevent.sleep(0)
except KeyboardInterrupt:
self.running = False
finally:
for device in devices:
device.close()
gevent.kill(crypto, KeyboardInterrupt)
gevent.kill(console_updater, KeyboardInterrupt)
def handler(self, data):
"""
Receives packets from headset for Windows. Sends them to a Queue to be processed
by the crypto greenlet.
"""
assert data[0] == 0
tasks.put_nowait("".join(map(chr, data[1:])))
self.packets_received += 1
return True
def setup_posix(self):
"""
Setup for headset on the Linux platform.
Receives packets from headset and sends them to a Queue to be processed
by the crypto greenlet.
"""
_os_decryption = False
if os.path.exists("/dev/eeg/raw"):
# The decryption is handled by the Linux epoc daemon. We don't need to handle it.
_os_decryption = True
hidraw = open("/dev/eeg/raw")
else:
serial, hidraw_filename = get_linux_setup()
self.serial_number = serial
if os.path.exists("/dev/" + hidraw_filename):
hidraw = open("/dev/" + hidraw_filename)
else:
hidraw = open("/dev/hidraw4")
crypto = gevent.spawn(self.setup_crypto, self.serial_number)
console_updater = gevent.spawn(self.update_console)
while self.running:
try:
data = hidraw.read(32)
if data != "":
if _os_decryption:
self.packets.put_nowait(EmotivPacket(data))
else:
# Queue it!
self.packets_received += 1
tasks.put_nowait(data)
gevent.sleep(0)
else:
# No new data from the device; yield
# We cannot sleep(0) here because that would go 100% CPU if both queues are empty
gevent.sleep(DEVICE_POLL_INTERVAL)
except KeyboardInterrupt:
self.running = False
hidraw.close()
if not _os_decryption:
gevent.kill(crypto, KeyboardInterrupt)
gevent.kill(console_updater, KeyboardInterrupt)
def setup_darwin(self):
"""
Setup for headset on the OS X platform.
Receives packets from headset and sends them to a Queue to be processed
by the crypto greenlet.
"""
_os_decryption = False
# Change these values to the hex equivalent from the output of hid_enumerate. If they are incorrect.
# Current values = VendorID: 8609 ProductID: 1
# You can see yours by opening /Applications/Utilities/System Information and navigating to
# Hardware -> USB -> Receiver Dongle L01
try:
hidraw = hid.device(0x21A1, 0x0001)
except IOError:
print "You need to specify correct ProductID and VendorID in setup_darwin()"
exit()
if self.serial_number == "":
print "Serial number needs to be specified manually in __init__()."
raise ValueError
crypto = gevent.spawn(self.setup_crypto, self.serial_number)
console_updater = gevent.spawn(self.update_console)
zero = 0
while self.running:
try:
# Doesn't seem to matter how big we make the buffer 32 returned every time, 33 for other platforms
data = hidraw.read(34)
if len(data) == 32:
# Most of the time the 0 is truncated? That's ok we'll add it...
data = [zero] + data
if data != "":
if _os_decryption:
self.packets.put_nowait(EmotivPacket(data))
else:
# Queue it!
tasks.put_nowait("".join(map(chr, data[1:])))
self.packets_received += 1
gevent.sleep(0)
else:
# No new data from the device; yield
# We cannot sleep(0) here because that would go 100% CPU if both queues are empty.
gevent.sleep(DEVICE_POLL_INTERVAL)
except KeyboardInterrupt:
self.running = False
hidraw.close()
gevent.kill(crypto, KeyboardInterrupt)
gevent.kill(console_updater, KeyboardInterrupt)
def setup_crypto(self, sn):
"""
Performs decryption of packets received. Stores decrypted packets in a Queue for use.
"""
if is_old_model(sn):
self.old_model = True
print self.old_model
k = ["\0"] * 16
k[0] = sn[-1]
k[1] = "\0"
k[2] = sn[-2]
if self.is_research:
k[3] = "H"
k[4] = sn[-1]
k[5] = "\0"
k[6] = sn[-2]
k[7] = "T"
k[8] = sn[-3]
k[9] = "\x10"
k[10] = sn[-4]
k[11] = "B"
else:
k[3] = "T"
k[4] = sn[-3]
k[5] = "\x10"
k[6] = sn[-4]
k[7] = "B"
k[8] = sn[-1]
k[9] = "\0"
k[10] = sn[-2]
k[11] = "H"
k[12] = sn[-3]
k[13] = "\0"
k[14] = sn[-4]
k[15] = "P"
key = "".join(k)
iv = Random.new().read(AES.block_size)
cipher = AES.new(key, AES.MODE_ECB, iv)
for i in k:
print "0x%.02x " % (ord(i))
while self.running:
while not tasks.empty():
task = tasks.get()
try:
data = cipher.decrypt(task[:16]) + cipher.decrypt(task[16:])
self.packets.put_nowait(EmotivPacket(data, self.sensors, self.old_model))
self.packets_processed += 1
except:
pass
gevent.sleep(0)
gevent.sleep(0)
def dequeue(self):
"""
Returns an EmotivPacket popped off the Queue.
"""
try:
return self.packets.get()
except Exception, e:
print e
_q = Queue()
_bs = BoundedSemaphore(10) #数据库并发最大10
_count = 0
def write_db(req):
'模拟写DB，随机休眠1-3秒，执行完毕释放信号'
global _count
_count += 1
print 'write %s to db' % req
sleep(random.randint(1,3))
_bs.release()
_count -= 1
def read_queue():
'读队列协程，信号量不为空切队列不为空时执行'
while True:
_bs.acquire()
req = _q.get()
spawn(write_db, req)
def info():
while True:
print 'qsize:%s,concurrent:%s' % (_q.qsize(),_count)
sleep(1)
if __name__ == '__main__':
spawn(read_queue)
spawn(info)
#模拟每秒接受5-30个请求
while True:
for i in xrange(random.randint(5,30)):
_q.put_nowait(random.randint(0,1000))
sleep(1)
class Emotiv(object):
def __init__(self, displayOutput=False, headsetId=0, research_headset=True):
self._goOn = True
self.packets = Queue()
self.packetsReceived = 0
self.packetsProcessed = 0
self.battery = 0
self.displayOutput = displayOutput
self.headsetId = headsetId
self.research_headset = research_headset
self.sensors = {
'F3': {'value': 0, 'quality': 0},
'FC6': {'value': 0, 'quality': 0},
'P7': {'value': 0, 'quality': 0},
'T8': {'value': 0, 'quality': 0},
'F7': {'value': 0, 'quality': 0},
'F8': {'value': 0, 'quality': 0},
'T7': {'value': 0, 'quality': 0},
'P8': {'value': 0, 'quality': 0},
'AF4': {'value': 0, 'quality': 0},
'F4': {'value': 0, 'quality': 0},
'AF3': {'value': 0, 'quality': 0},
'O2': {'value': 0, 'quality': 0},
'O1': {'value': 0, 'quality': 0},
'FC5': {'value': 0, 'quality': 0},
'X': {'value': 0, 'quality': 0},
'Y': {'value': 0, 'quality': 0},
'Unknown': {'value': 0, 'quality': 0}
}
def setup(self, headsetId=0):
if windows:
self.setupWin()
else:
self.setupPosix()
def updateStdout(self):
while self._goOn:
if self.displayOutput:
if windows:
os.system('cls')
else:
os.system('clear')
print "Packets Received: %s Packets Processed: %s" % (self.packetsReceived, self.packetsProcessed)
print('\n'.join("%s Reading: %s Strength: %s" % (k[1], self.sensors[k[1]]['value'],self.sensors[k[1]]['quality']) for k in enumerate(self.sensors)))
print "Battery: %i" % g_battery
gevent.sleep(1)
def getLinuxSetup(self):
rawinputs = []
for filename in os.listdir("/sys/class/hidraw"):
realInputPath = check_output(["realpath", "/sys/class/hidraw/" + filename])
sPaths = realInputPath.split('/')
s = len(sPaths)
s = s - 4
i = 0
path = ""
while s > i:
path = path + sPaths[i] + "/"
i += 1
rawinputs.append([path, filename])
hiddevices = []
# TODO: Add support for multiple USB sticks? make a bit more elegant
for input in rawinputs:
try:
with open(input[0] + "/manufacturer", 'r') as f:
manufacturer = f.readline()
f.close()
if ("Emotiv Systems Inc." in manufacturer) or ("Emotiv Systems Pty Ltd" in manufacturer) :
with open(input[0] + "/serial", 'r') as f:
serial = f.readline().strip()
f.close()
print "Serial: " + serial + " Device: " + input[1]
# Great we found it. But we need to use the second one...
hidraw = input[1]
id_hidraw = int(hidraw[-1])
# The dev headset might use the first device, or maybe if more than one are connected they might.
id_hidraw += 1
hidraw = "hidraw" + id_hidraw.__str__()
print "Serial: " + serial + " Device: " + hidraw + " (Active)"
return [serial, hidraw, ]
except IOError as e:
print "Couldn't open file: %s" % e
def setupWin(self):
devices = []
try:
for device in hid.find_all_hid_devices():
if device.vendor_id != 0x21A1:
continue
if device.product_name == 'Brain Waves':
devices.append(device)
device.open()
self.serialNum = device.serial_number
device.set_raw_data_handler(self.handler)
elif device.product_name == 'EPOC BCI':
devices.append(device)
device.open()
self.serialNum = device.serial_number
device.set_raw_data_handler(self.handler)
elif device.product_name == '00000000000':
devices.append(device)
device.open()
self.serialNum = device.serial_number
device.set_raw_data_handler(self.handler)
gevent.spawn(self.setupCrypto, self.serialNum)
gevent.spawn(self.updateStdout)
while self._goOn:
try:
gevent.sleep(0)
except KeyboardInterrupt:
self._goOn = False
for device in devices:
device.close()
finally:
for device in devices:
device.close()
def handler(self, data):
assert data[0] == 0
tasks.put_nowait(''.join(map(chr, data[1:])))
self.packetsReceived += 1
return True
def setupPosix(self):
_os_decryption = False
if os.path.exists('/dev/eeg/raw'):
# The decrpytion is handled by the Linux epoc daemon. We don't need to handle it there.
_os_decryption = True
self.hidraw = open("/dev/eeg/raw")
else:
setup = self.getLinuxSetup()
self.serialNum = setup[0]
if os.path.exists("/dev/" + setup[1]):
self.hidraw = open("/dev/" + setup[1])
else:
self.hidraw = open("/dev/hidraw4")
gevent.spawn(self.setupCrypto, self.serialNum)
gevent.spawn(self.updateStdout)
while self._goOn:
try:
data = self.hidraw.read(32)
if data != "":
if _os_decryption:
self.packets.put_nowait(EmotivPacket(data))
else:
# Queue it!
self.packetsReceived += 1
tasks.put_nowait(data)
gevent.sleep(0)
except KeyboardInterrupt:
self._goOn = False
return True
def setupCrypto(self, sn):
type = 0 # feature[5]
type &= 0xF
type = 0
# I believe type == True is for the Dev headset, I'm not using that. That's the point of this library in the first place I thought.
k = ['\0'] * 16
k[0] = sn[-1]
k[1] = '\0'
k[2] = sn[-2]
if type:
k[3] = 'H'
k[4] = sn[-1]
k[5] = '\0'
k[6] = sn[-2]
k[7] = 'T'
k[8] = sn[-3]
k[9] = '\x10'
k[10] = sn[-4]
k[11] = 'B'
else:
k[3] = 'T'
k[4] = sn[-3]
k[5] = '\x10'
k[6] = sn[-4]
k[7] = 'B'
k[8] = sn[-1]
k[9] = '\0'
k[10] = sn[-2]
k[11] = 'H'
k[12] = sn[-3]
k[13] = '\0'
k[14] = sn[-4]
k[15] = 'P'
# It doesn't make sense to have more than one greenlet handling this as data needs to be in order anyhow. I guess you could assign an ID or something
# to each packet but that seems like a waste also or is it? The ID might be useful if your using multiple headsets or usb sticks.
key = ''.join(k)
iv = Random.new().read(AES.block_size)
cipher = AES.new(key, AES.MODE_ECB, iv)
for i in k: print "0x%.02x " % (ord(i))
while self._goOn:
while not tasks.empty():
task = tasks.get()
data = cipher.decrypt(task[:16]) + cipher.decrypt(task[16:])
self.lastPacket = EmotivPacket(data, self.sensors)
self.packets.put_nowait(self.lastPacket)
self.packetsProcessed += 1
gevent.sleep(0)
gevent.sleep(0)
def dequeue(self):
try:
return self.packets.get()
except Exception, e:
print e
