def handle_agent(sock): agent = None ok = False while True: try: raw = sock.recv(int(1e8)).decode() #data = json.loads(raw) data = json_decoder.decode(raw) action = data['action'] r = handle_agent_action(sock, agent, action, data) if 'agent' in r: agent = r['agent'] if action != 'reply': r['action'] = 'reply' sock.send(json.dumps(r).encode()) if action == 'register': output_status() ok = True except (ConnectionResetError, ConnectionAbortedError) as e: import traceback traceback.print_exc() if agent: remove_agent(agent) del agents[agent] output_status() break except json.decoder.JSONDecodeError as e: prt('data error from agent', agent, type(e), e) prt('raw data:', raw) if not ok: break ok = False sock.close()
def handle_action(websocket, action, data): if action == 'register': company = data['company'] department = '%s-%s' % (company, data['department']) group = '%s-%s' % (department, data['group']) timestamp,_,digest = data['digest'].partition('|') if group in tokens: token = tokens[group] elif department in tokens: token = tokens[department] else: token = tokens[company] assert hashlib.md5((timestamp+'|'+token).encode()).hexdigest() == digest websocket2feeder[websocket] = dict(company=company, department=department, group=group, agent2data={}) return dict(status='ok') elif action == 'add-agent': agent = data['agent'] agent_data = dict(inputs=data['inputs'], cleanses=data['cleanses'], outputs=data['outputs']) websocket2feeder[websocket]['agent2data'][agent] = agent_data return dict(status='ok') elif action == 'remove-agent': agent = data['agent'] agent2data = websocket2feeder[websocket]['agent2data'] if agent in agent2data: del agent2data[agent] return dict(status='ok') elif action == 'reply': prt('GOT REPLY:', action, data) job = data['job'] job2reply[job] = (websocket,data) return dict(status='ok') else: prt('BAD DATA FROM FEEDER:', action, data) return dict(status='error')
def get_job(job): result = dict(status='pending') if job in job2reply: websocket,reply = job2reply[job] prt('getting job reply:', reply) result = dict(reply) del result['action'] del result['job'] if request.values.get('dom') != None: del result['data'] if reply['reference-action'] == 'find': agent_hits = reply['data'] cleanses = {} allow_cleanse = True if websocket in websocket2feeder: feeder_data = websocket2feeder[websocket] feeder_datas = get_company_feeders(feeder_data['company']).values() for feeder_data in feeder_datas: for agent,init_data in feeder_data['agent2data'].items(): agent_hit_user = agent_hits[agent]['data'] allow_cleanse &= (len(agent_hit_user) == 1) for user in agent_hit_user: cleanses.update({k:user[k] for k in init_data['cleanses'] if k in user}) if not cleanses: allow_cleanse = False result['dom'] = render_template('search-result.html', agents_hits=agent_hits, cleanses=cleanses, allow_cleanse=allow_cleanse) else: line_cnt = 0 for agent, agent_data in reply['data'].items(): if agent_data['status'] == 'ok': line_cnt += agent_data['data']['lines'] result['dom'] = render_template('cleanse-result.html', line_cnt=line_cnt) return jsonify(result)
def _provider_add_addr(web, prv_name, prv_addr, kms_addr_spec, prv_agt): pnbox = web.find('elmtNm') pnbox.clear() web.sendkeys(pnbox, prv_name) stale = web.check_stale(_PRV_NAME_IN_LIST.format(prv_name), SW.CSS) # maybe target PRV is already listed web.click('searchBtn') # Search for target PRV will obsolete old element if stale: web.wait_detach( stale ) # Make sure to pick a fresh element after the old one got detached try: web.clickp(_PRV_NAME_IN_LIST.format(prv_name), SW.CSS) addrs_exist = [] _walkthrough_kms_addrs( web, lambda re, ka, pa: addrs_exist.append( (ka, pa)) is 'Return False') pinf('add provider addresses for Provider[{}]...', prv_addr) for i, kms_addr in enumerate(Conf.expand_kms_addr(kms_addr_spec), start=1): if (kms_addr, prv_addr) in addrs_exist: perr(' {}: {} exist/skipped', i, kms_addr) continue web.click('protIntfInsert') web.sendkeys(web.find('kmsProtocolId'), kms_addr) web.sendkeys(web.find('provProtocolId'), prv_addr) web.click('app') web.click(f'#app option[value="{prv_agt}"]', SW.CSS) web.click('protIntfSaveBtn') row_add = _walkthrough_kms_addrs( web, lambda re, ka, pa: ka == kms_addr and pa == prv_addr) web.click(row_add) prt(' {}: {}/{} added', i, kms_addr, prv_agt) else: # save to DB if 'row_add' in dir(): web.click('btnPrvdrSave') # save all prt(' committed') # accept modal dialogue web.click('div.modal-content button.bootbox-accept', SW.CSS) return perr(' aborted !') except TimeoutException: perr('provider: {} not found', prv_name) except StaleElementReferenceException as e: pause(2) raise e
def walkthrough_callback(re, ka, pa): nonlocal bDeleteDirty if pa == prv_addr and ka in kaddrs2del: chkbox = re.find_element_by_tag_name('input') web.click(chkbox) # check row to delete prt(' {} : {} checked ', ka, pa) kaddrs2del.remove(ka) bDeleteDirty = True return False
def _provider_del_addr(web, prv_name, prv_addr, kms_addr_spec, prv_agt): pnbox = web.find('elmtNm') pnbox.clear() web.sendkeys(pnbox, prv_name) stale = web.check_stale(_PRV_NAME_IN_LIST.format(prv_name), SW.CSS) # maybe target PRV is already listed web.click('searchBtn') # Search for target PRV will obsolete old element if stale: web.wait_detach( stale ) # Make sure to pick a fresh element after the old one got detached try: web.clickp(_PRV_NAME_IN_LIST.format(prv_name), SW.CSS) pinf('delete provider addresses for Provider[{}]...', prv_addr) kaddrs2del = list(Conf.expand_kms_addr(kms_addr_spec)) bDeleteDirty = False def walkthrough_callback(re, ka, pa): nonlocal bDeleteDirty if pa == prv_addr and ka in kaddrs2del: chkbox = re.find_element_by_tag_name('input') web.click(chkbox) # check row to delete prt(' {} : {} checked ', ka, pa) kaddrs2del.remove(ka) bDeleteDirty = True return False _walkthrough_kms_addrs(web, walkthrough_callback) if len(kaddrs2del) > 0: for ka in kaddrs2del: perr(' {} : {} not exist/skipped', ka, prv_addr) # commit if bDeleteDirty: web.click('protIntfDelete') web.click('protIntfDelOK') prt(' deleted') web.click('btnPrvdrSave') prt(' committed') # accept modal dialogue web.click('div.modal-content button.bootbox-accept', SW.CSS) except StaleElementReferenceException as e: pause(2) raise e
def _walkthrough_kms_addrs(web, callee=None): try: for r in web.finds(_PRV_ADDRS_TRS): kms_addr, prv_addr = r.find_elements_by_css_selector( _KMS_PRV_ADDR_TDS) if not callee: prt(' - prv addr = {}:{}', kms_addr.text, prv_addr.text) else: if callee(r, kms_addr.text, prv_addr.text): return r else: return None except TimeoutException: return None
async def handle(websocket, path): try: async for message in websocket: data = json.loads(message) prt('data from feeder:', data) action = data['action'] r = handle_action(websocket, action, data) if action != 'reply': r['action'] = 'reply' r['reference-action'] = action await websocket.send(json.dumps(r)) except Exception as e: prt(type(e), e) finally: if websocket in websocket2feeder: del websocket2feeder[websocket]
def index(): nxt = request.values.get('next', '/search') form = LoginForm(request.form) if request.method == 'POST' and form.validate(): username = form.username.data password = form.password.data domain = username.partition('@')[2].split('.')[-2] fullname, groups = login(domain, username, password) user = load_user(username) user.domain = domain user.fullname = fullname user.groups = groups login_user(user) assert current_user.fullname == fullname prt('Logged in successfully.') return redirect(nxt) return render_template('index.html', form=form, next=nxt)
def test_gen(gen, n=5): Xtr, Ytr = next(gen) for i in range(n): assert Xtr[i, maxlend] == eos x = Xtr[i, :maxlend] y = Xtr[i, maxlend:] yy = Ytr[i, :] yy = np.where(yy)[1] prt('L', idx2word, yy) prt('H', idx2word, y) if maxlend: prt('D', idx2word, x)
def send_req(action, fields, req_data): prt(action, 'data:', req_data) sent_to = [] for agent, agent_data in agents.items(): agent_fields = set(agent_data[fields]) prt('agent data: ', agent_data) query_data = { key: value for key, value in req_data.items() if value and key in agent_fields } prt('query data:', query_data) if not query_data: continue send_data = dict(action=action, id=req_data['id'], query=query_data) sock = agent_data['socket'] sock.send(json.dumps(send_data).encode()) sent_to += [agent] return sent_to
def output_status(): prt('status: %s' % ('ok' if connected2hoarder else 'not connected')) prt('agents: %i' % len(agents))
def service_master(uri): while True: raw = '' try: global connected2hoarder, ws connected2hoarder = False prt('connecting to', uri) ws = websock = websocket.create_connection(uri) prt('after connect') timestamp = datetime.now().isoformat() digest = hashlib.md5( (timestamp + '|' + token).encode()).hexdigest() digest = timestamp + '|' + digest data = json.dumps( dict(action='register', company=company, department=department, group=group, digest=digest)) websock.send(data) for agent in agents: add_agent(agent) while True: raw = websock.recv() data = json.loads(raw) connected2hoarder = True prt('got hoarder data:', data) action = data['action'] handle_hoarder_action(websock, action, data) except Exception as e: import os env = set(k.lower() for k in os.environ.keys()) if 'http_proxy' in env or 'https_proxy' in env: prt('Turn off proxy?') else: prt(type(e), e) prt(raw)
def handle_hoarder_action(websock, action, data): if action in ('find', 'cleanse'): global agent_job2queue req_data = dict(data['data']) assert req_data job = str(uuid.uuid4()) req_data['id'] = job #agent_job2queue[job] = asyncio.Queue() agent_job2queue[job] = queue.Queue() fields = 'inputs' if action == 'find' else 'cleanses' agents = send_req(action, fields, req_data) agents_data = {agent: {} for agent in agents} recv_cnt = len(agents) prt('%s waiting for %i answers from %i agents...' % (action, recv_cnt, len(agents))) for cnt in range(recv_cnt): # explicit about count r = agent_job2queue[job].get() prt(action, 'got a reply:', r) agent = r['agent'] agents_data[agent].update(r) del agents_data[agent]['agent'] reply = dict(status=('ok' if cnt == recv_cnt - 1 else 'ok-partial'), action='reply', job=data['job'], data=agents_data) reply['reference-action'] = action prt(action, 'partial reply:', reply) websock.send(json.dumps(reply)) prt(action, 'done') elif action == 'login': username = data['username'] password = data['password'] domain = data['domain'] fullname, groups = ldap_login.login(domain, username, password) prt('login:'******'ok', action='reply', job=data['job'], fullname=fullname, groups=groups) reply['reference-action'] = action websock.send(json.dumps(reply)) elif action == 'reply': pass else: prt(data) assert False
def forward(agent, job, data): data['agent'] = agent prt('forwarding', data) #ws_loop.call_soon_threadsafe(lambda: asyncio.ensure_future(agent_job2queue[job].put(data))) agent_job2queue[job].put(data)
def send2feeder(websocket, obj): prt('sending to feeder:', obj) data = json.dumps(obj) main_event_loop.call_soon_threadsafe(lambda: asyncio.ensure_future(websocket.send(data)))