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)