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)))
