async def mark_profiles_as_not_running(): all_profiles = await db_objects.execute( C2Profile.select().where(C2Profile.running == True)) for profile in all_profiles: profile.running = False await db_objects.update(profile) return
async def get_all_c2profiles_for_current_operation(request, user): if user['current_operation'] != "": try: operation = await db_objects.get(Operation, name=user['current_operation']) all_profiles = await db_objects.execute( C2Profile.select().where(C2Profile.operation == operation)) profiles = await db_objects.execute( PayloadTypeC2Profile.select( PayloadTypeC2Profile, C2Profile, PayloadType).join(C2Profile).switch( PayloadTypeC2Profile).join(PayloadType)) results = [] inter = {} for p in all_profiles: inter[p.name] = p.to_json() inter[p.name]['ptype'] = [ ] # start to keep track of which payload types this profile supports for p in profiles: if p.c2_profile.operation == operation: inter[p.c2_profile.name]['ptype'].append( p.payload_type.ptype) for k in inter.keys(): # make an array of dictionaries results.append(inter[k]) return json(results) except Exception as e: print(e) return json([""]) else: return json([""])
async def ws_c2profile_current_operation(request, ws, user): try: async with aiopg.create_pool(apfell.config['DB_POOL_CONNECT_STRING']) as pool: async with pool.acquire() as conn: async with conn.cursor() as cur: await cur.execute('LISTEN "newc2profile";') # BEFORE WE START GETTING NEW THINGS, UPDATE WITH ALL OF THE OLD DATA if user['current_operation'] != "": operation = await db_objects.get(Operation, name=user['current_operation']) profiles = await db_objects.execute(C2Profile.select().where(C2Profile.operation == operation).order_by(C2Profile.id)) for p in profiles: await ws.send(js.dumps(p.to_json())) await ws.send("") # now pull off any new payloads we got queued up while processing old data while True: try: msg = conn.notifies.get_nowait() id = (msg.payload) p = await db_objects.get(C2Profile, id=id) if p.operation == operation: await ws.send(js.dumps(p.to_json())) except asyncio.QueueEmpty as e: await asyncio.sleep(1) await ws.send("") # this is our test to see if the client is still there continue except Exception as e: print(e) return finally: pool.close()
async def get_all_c2profiles(request, user): # this syntax is atrocious for getting a pretty version of the results from a many-to-many join table) all_profiles = await db_objects.execute(C2Profile.select()) profiles = await db_objects.execute( PayloadTypeC2Profile.select(PayloadTypeC2Profile, C2Profile, PayloadType).join(C2Profile).switch( PayloadTypeC2Profile).join(PayloadType) ) results = [] inter = {} for p in all_profiles: # only show profiles for operations the user is part of # overall admins can see all operations if p.operation.name in user['operations'] or user['admin']: inter[p.name] = p.to_json() # create an empty array for ptypes that we'll populate in the next for loop inter[p.name]['ptype'] = [] for p in profiles: if p.c2_profile.name in inter: inter[p.c2_profile.name]['ptype'].append(p.payload_type.ptype) for k in inter.keys(): results.append(inter[k]) return json(results)