-
Notifications
You must be signed in to change notification settings - Fork 57
/
engine.py
334 lines (259 loc) · 8.22 KB
/
engine.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
"""
For communication with Cobalt Strike
"""
# TODO better argument checking on aggressor functions
import json
import re
import sys
import traceback
import os
#sys.path.insert(0, os.path.realpath(os.path.dirname(__file__)) + '/..')
import pycobalt.utils as utils
import pycobalt.callbacks as callbacks
import pycobalt.serialization as serialization
_in_pipe = None
_out_pipe = None
_debug_on = False
def _init_pipes():
"""
Configure input and output pipes. At the moment we use stdin/out/err. This
just makes configuring scripts a bit easier. It would be pretty easy to use
a couple of fifos but then we have to pass them to the Python script.
Passing them on argv seems kind of dirty too.
"""
global _in_pipe
global _out_pipe
_in_pipe = sys.stdin
_out_pipe = sys.stdout
def enable_debug():
"""
Enable debug messages on the Python side
To enable the Aggressor debug messages run `python-debug` in the Script
Console or set `$pycobalt_debug_on = true` in your Aggressor script.
"""
global _debug_on
_debug_on = True
debug('enabled debug')
def disable_debug():
"""
Disable debug messages
"""
global _debug_on
debug('disabling debug')
_debug_on = False
def debug(line):
"""
Write script console debug message
:param line: Line to write
"""
global _debug_on
if _debug_on:
write('debug', str(line))
def handle_exception_softly(exc):
"""
Print an exception to the script console
:param exc: Exception to print
"""
try:
raise exc
except Exception as e:
error('Exception: {}\n'.format(str(e)))
error('Traceback: {}'.format(traceback.format_exc()))
def write(message_type, message=''):
"""
Write a message to Cobalt Strike. Message can be anything serializable by
`serialization.py`. This includes primitives, bytes, lists, dicts, tuples,
and callbacks (automatically registered).
:param message_type: Type/label of message
:param message: Message contents
"""
global _out_pipe
wrapper = {
'name': message_type,
'message': message,
}
serialized = serialization.serialized(wrapper)
_out_pipe.write(serialized + "\n")
_out_pipe.flush()
def handle_message(name, message):
"""
Handle a received message according to its name
:param name: Name/type/label of message
:param message: Message body
"""
#debug('handling message of type {}: {}'.format(name, message))
if name == 'callback':
# dispatch callback
callback_name = message['name']
callback_args = message['args'] if 'args' in message else []
if 'sync' in message and message['sync'] and 'id' in message:
return_message = callbacks.call(callback_name, callback_args, return_id=message['id'])
write('return', return_message)
else:
return_message = callbacks.call(callback_name, callback_args)
elif name == 'eval':
# eval python code
eval(message)
elif name == 'debug':
# set debug mode
if message is True:
enable_debug()
else:
disable_debug()
elif name == 'stop':
# stop script
stop()
elif not name:
error('Error reading pipe: {}'.format(str(message)))
handle_exception_softly(message)
else:
error('Received unhandled or out-of-order message type: {} {}'.format(name, str(message)))
def parse_line(line):
"""
Parse a serialized input line for passing to `engine.handle_message`.
:param line: Line to parse. Should look like {'name':<name>, 'message':<message>}
:return: Tuple containing 'name' and 'message'
"""
try:
# remove shitty unicode
line = line.encode('utf-8', 'ignore').decode()
line = line.strip()
wrapper = json.loads(line, strict=False)
name = wrapper['name']
if 'message' in wrapper:
message = wrapper['message']
else:
message = None
return name, message
except Exception as e:
return None, e
_has_forked = False
def fork():
"""
Tell Cobalt Strike to fork into a new thread.
Menu trees have to be registered before we fork into a new thread so this
is called in `engine.loop()` after the registration is finished.
"""
global _has_forked
if _has_forked:
raise RuntimeError('Tried to fork Cobalt Strike twice')
write('fork')
def read_pipe():
"""
read_pipe a message line
:return: Tuple containing message name and contents (as returned by
`parse_line`).
"""
global _in_pipe
return parse_line(next(_in_pipe))
def read_pipe_iter():
"""
read_pipe message lines
:return: Iterator with an item for each read_pipe/parsed line. Each item is the
same as the return value of `engine.read_pipe()`.
"""
global _in_pipe
for line in _in_pipe:
yield parse_line(line)
def loop(fork_first=True):
"""
Loop forever, handling messages. Does not return until the pipe closes.
Exceptions are printed to the script console.
:param fork_first: Whether to call `fork()` first.
"""
global _has_forked
if fork_first and not _has_forked:
fork()
for name, message in read_pipe_iter():
try:
if name:
handle_message(name, message)
else:
error('Received invalid message: {}'.format(message))
except Exception as e:
handle_exception_softly(e)
def stop():
"""
Stop the script (just exits the process)
"""
sys.exit()
def call(name, args=None, silent=False, fork=None, sync=True):
"""
Call a sleep/aggressor function. You should use the `aggressor.py` helpers
where possible.
:param name: Name of function to call
:param args: Arguments to pass to function
:param silent: Don't print tasking information (! operation) (only works
for some functions)
:param fork: Call in its own thread
:param sync: Wait for return value
:return: Return value of function if `sync` is True
"""
if args is None:
# no arguments
args = []
# serialize and register function callbacks if needed
if fork is None:
if callbacks.has_callback(args):
# when there's a callback involved we usually have to fork because the
# main script thread is busy reading from the script.
#debug("forcing fork for call to: {}".format(name))
fork = True
else:
fork = False
message = {
'name': name,
'args': args,
'silent': silent,
'fork': fork,
'sync': sync,
}
write('call', message)
if sync:
# read_pipe and handle messages until we get our return value
for name, message in read_pipe_iter():
if name == 'return':
# got it
return message
else:
try:
handle_message(name, message)
except Exception as e:
handle_exception_softly(e)
def eval(code):
"""
Eval aggressor code. Does not provide a return value.
:param code: Code to eval
"""
write('eval', code)
def menu(menu_items):
"""
Register a Cobalt Strike menu tree
:param menu_items: Menu tree as returned by the `gui.py` helpers.
"""
global _has_forked
if _has_forked:
raise RuntimeError('Tried to register a menu after forking. This crashes Cobalt Strike')
write('menu', menu_items)
def error(line):
"""
Write error notice
:param line: Line to write
"""
write('error', str(line))
def message(line):
"""
Write script console message. The Aggressor side will add a prefix to your
message. To print raw messages use `aggressor.print` and
`aggressor.println`.
:param line: Line to write
"""
write('message', str(line))
def delete(handle):
"""
Delete an object with its serialized handle. This just removed the global
reference. The object will stick around if it's referenced elsewhere.
:param handle: Handle of object to delete
"""
write('delete', handle)
_init_pipes()