class CustomShell: def __init__(self, services): self.log = Logger('terminal') self.shell_prompt = 'caldera> ' self.modes = dict(session=Session(services, self.log), agent=Agent(services, self.log), ability=Ability(services, self.log), adversary=Adversary(services, self.log), operation=Operation(services, self.log)) async def start_shell(self): await asyncio.sleep(1) while True: try: cmd = await ainput(self.shell_prompt) self.log.debug(cmd) mode = re.search(r'\((.*?)\)', self.shell_prompt) if cmd == 'help': await self._print_help() elif cmd.startswith('log'): await self._print_logs(int(cmd.split(' ')[1])) elif cmd in self.modes.keys(): self.shell_prompt = 'caldera (%s)> ' % cmd elif mode: await self.modes[mode.group(1)].execute(cmd) elif cmd == '': pass else: self.log.console( 'Bad command - are you in the right mode?', 'red') except Exception as e: self.log.console('Bad command: %s' % e, 'red') async def accept_sessions(self, reader, writer): address = writer.get_extra_info('peername') connection = writer.get_extra_info('socket') connection.setblocking(1) self.modes['session'].sessions.append(connection) self.modes['session'].addresses.append('%s:%s' % (address[0], address[1])) self.log.console('New session: %s:%s' % (address[0], address[1])) async def _print_help(self): print('HELP MENU:') print('-> help: show this help menu') print('-> logs [n]: view the last n-lines of each log file') print( 'Enter one of the following modes. Once inside, enter "info" to see available commands.' ) for cmd, v in self.modes.items(): print('-> %s' % cmd) @staticmethod async def _print_logs(n): for name in glob.iglob('.logs/*.log', recursive=False): with open(name, 'r') as f: print('***** %s ***** ' % name) lines = f.readlines() print(*lines[-n:])
class OperationService(OpControl): def __init__(self, data_svc, utility_svc, planner): super().__init__(data_svc.dao) self.data_svc = data_svc self.utility_svc = utility_svc self.loop = asyncio.get_event_loop() self.log = Logger('operation') planning_module = import_module(planner) self.planner = getattr(planning_module, 'LogicalPlanner')(self.data_svc, self.utility_svc, self.log) async def resume(self): for op in await self.data_svc.dao.get('core_operation'): if not op['finish']: self.loop.create_task(self.run(op['id'])) async def close_operation(self, op_id): self.log.console('Operation complete: %s' % op_id, 'blue') update = dict(finish=datetime.now().strftime('%Y-%m-%d %H:%M:%S')) await self.data_svc.dao.update('core_operation', key='id', value=op_id, data=update) async def run(self, op_id): self.log.console('Starting operation: %s' % op_id, 'blue') operation = await self.data_svc.explode_operation(dict(id=op_id)) try: for phase in operation[0]['adversary']['phases']: self.log.debug('Operation %s phase %s: started' % (op_id, phase)) await self.planner.execute(operation[0], phase) self.log.debug('Operation %s phase %s: completed' % (op_id, phase)) await self.data_svc.dao.update('core_operation', key='id', value=op_id, data=dict(phase=phase)) operation = await self.data_svc.explode_operation( dict(id=op_id)) if operation[0]['cleanup']: await self.cleanup(op_id) await self.close_operation(op_id) except Exception: traceback.print_exc() async def cleanup(self, op_id): self.log.console('Cleanup started for operation: %s' % op_id, 'blue') clean_commands = await self.data_svc.dao.get('core_cleanup', dict(op_id=op_id)) for c in reversed(clean_commands): link = dict(op_id=c['op_id'], host_id=c['agent_id'], ability_id=c['ability_id'], decide=datetime.now(), command=c['command'], score=0, jitter=1) await self.data_svc.create_link(link) await self.cleanup_operation(op_id)