-
Notifications
You must be signed in to change notification settings - Fork 57
/
aliases.py
143 lines (114 loc) · 5 KB
/
aliases.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
"""
For registering beacon console aliasess
Regular example:
def test_alias(args):
print(args)
aliases.register('test_alias', test_alias)
Decorator example:
@aliases.alias('test_alias')
def test_alias(args):
print(args)
"""
import pycobalt.aggressor as aggressor
import pycobalt.callbacks as callbacks
import pycobalt.engine as engine
import pycobalt.utils as utils
_default_quote_replacement = None
def set_quote_replacement(replacement):
"""
Set the default `quote_replacement` value. Passing `quote_replacement=` to
`register()` or `@alias()` overrides this.
See `register()` for more information.
:param replacement: Quote replacement string
"""
global _default_quote_replacement
_default_quote_replacement = replacement
def register(name, callback, short_help=None, long_help=None,
quote_replacement=None):
"""
Register an alias
Regarding the `quote_replacement` argument: Cobalt Strike's Beacon console
uses double quotes to enclose arguments with spaces in them. There's no way
to escape double quotes within those quotes though. Set `quote_replacement`
to a string and PyCobalt will replace it with " in each argument. Call
`aliases.set_quote_replacement(<string>)` to change the default quote
replacement behavior.
:param name: Name of alias
:param callback: Callback for alias
:param short_help: Short version of help, shown when running `help` with no
arguments
:param long_help: Long version of help, shown when running `help <alias>`.
By default this is generated based on the short help and syntax of the
Python callback.
:param quote_replacement: Replace this string with " in each
argument.
"""
# this is a workaround for a famously stupid python issue where keyword
# arguments are not passed to closures correctly.
quote_replacement_ = quote_replacement
def alias_callback(*args):
# first argument is bid
bid = int(args[0])
# see above
quote_replacement = quote_replacement_
# check arguments
if not utils.check_args(callback, args):
syntax = '{} {}'.format(name, utils.signature_command(callback, trim=1))
aggressor.berror(bid, "Syntax: " + syntax)
engine.error("Invalid number of arguments passed to alias '{}'. Syntax: {}".format(name, syntax))
return
# handle the quote replacement character
if not quote_replacement:
global _default_quote_replacement
quote_replacement = _default_quote_replacement
if quote_replacement:
args = [arg.replace(quote_replacement, '"') for arg in args]
try:
# run the alias callback
#engine.debug('calling callback for alias {}'.format(name))
callback(*args)
except Exception as e:
# print exception summaries to the beacon log. raise the
# Exception again so the full traceback can get printed to the
# Script Console
aggressor.berror(bid,
"Caught Python exception while executing alias '{}': {}\n See Script Console for more details.".format(name, str(e)))
raise e
callbacks.register(alias_callback, prefix='alias_{}'.format(name))
aggressor.alias(name, alias_callback, sync=False)
# by default short_help is just 'Custom Python command'
if not short_help:
short_help = 'Custom Python command'
# by default long_help is short_help
if not long_help:
long_help = short_help
# tack the syntax on the long_help, even if the user set their own
long_help += '\n\nSyntax: {} {}'.format(name, utils.signature_command(callback, trim=1))
aggressor.beacon_command_register(name, short_help, long_help, sync=False)
class alias:
"""
Decorator for alias registration
"""
def __init__(self, name, short_help=None, long_help=None,
quote_replacement=None):
"""
:param name: Name of alias
:param short_help: Short version of help, shown when running `help` with no
arguments
:param long_help: Long version of help, shown when running `help <alias>`.
By default this is generated based on the short help and syntax of the
Python callback.
:param quote_replacement: Replace this string with " in each
argument (see notes in the `register()`
function)
"""
self.name = name
self.short_help = short_help
self.long_help = long_help
self.quote_replacement = quote_replacement
def __call__(self, func):
self.func = func
register(self.name, self.func,
short_help=self.short_help,
long_help=self.long_help,
quote_replacement=self.quote_replacement)