def register(name, callback):
"""
Register a command
:param name: Name of command
:param callback: Callback for command
"""
def command_callback(*args):
engine.debug('calling callback for command {}'.format(name))
callback(*args)
callbacks.register(command_callback, prefix='command_{}'.format(name))
aggressor.command(name, command_callback)
def _serialize_special(item):
"""
Serialize and register callbacks and other special objects
"""
if isinstance(item, list) or isinstance(item, tuple):
# recurse lists
new_list = []
for child in item:
new_list.append(_serialize_special(child))
return new_list
elif isinstance(item, dict):
# recurse dicts
new_dict = {}
for key, value in item.items():
new_dict[key] = _serialize_special(value)
return new_dict
elif callable(item):
# serialize callbacks
func_name = callbacks.name(item)
if not func_name:
func_name = callbacks.register(item)
return _serialize_callback_prefix + func_name
elif isinstance(item, bytes):
# serialize bytes
return _serialize_bytes_prefix + base64.b64encode(item).decode()
else:
# item is json-serializable
return item
def register_modifier(name, callback, known_only=True):
"""
Register a modifier callback.
:param name: Name of modifier (case-insensitive)
:param callback: Modifier callback
:param known_only: Only allow known modifiers
:return: Name of registered callback
"""
name = name.upper()
def modifier_callback(*args):
engine.debug('calling callback for modifier {}'.format(name))
try:
return callback(*args)
except Exception as exc:
engine.handle_exception_softly(exc)
return '[!] An Exception occurred in the {} output modifier. See Script Console for more details.'.format(
name)
if known_only and not is_known_modifier(name):
raise RuntimeError(
'Tried to register an unknown modifier: {name}. Try console.modifier("{name}", known_only=False).'
.format(name=name))
callback_name = callbacks.register(modifier_callback,
prefix='modifier_{}'.format(name))
engine.write('set', {'name': name, 'callback': modifier_callback})
return callback_name
def register(name, callback, short_help=None, long_help=None):
"""
Register an alias
: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.
"""
def alias_callback(*args):
bid = int(args[0])
if utils.check_args(callback, args):
try:
engine.debug('calling callback for alias {}'.format(name))
callback(*args)
except Exception as e:
aggressor.berror(
bid,
"Caught Python exception while executing alias '{}': {}\n See Script Console for more details."
.format(name, str(e)))
raise e
else:
syntax = '{}{}'.format(name, utils.signature(callback, trim=1))
aggressor.berror(bid, "Syntax: " + syntax)
engine.error(
"Invalid number of arguments passed to alias '{}'. Syntax: {}".
format(name, syntax))
callbacks.register(alias_callback, prefix='alias_{}'.format(name))
aggressor.alias(name, alias_callback)
# register help info
if not long_help:
long_help = ''
if short_help:
long_help += short_help + '\n\n'
long_help += 'Syntax: {}{}'.format(name,
utils.signature(callback, trim=1))
if not short_help:
short_help = 'Custom python command'
aggressor.beacon_command_register(name, short_help, long_help)
def register(name, callback, quote_replacement=None):
"""
Register a command
Regarding the `quote_replacement` argument: Cobalt Strike's Script 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.
:param name: Name of command
:param callback: Callback for command
: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 command_callback(*args):
# see above
quote_replacement = quote_replacement_
# check arguments
if not utils.check_args(callback, args):
syntax = '{} {}'.format(name, utils.signature_command(callback))
engine.error("Syntax: " + 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]
#engine.debug('calling callback for command {}'.format(name))
callback(*args)
callbacks.register(command_callback, prefix='command_{}'.format(name))
aggressor.command(name, command_callback, sync=False)
def register(name, callback, official_only=True):
"""
Register an event callback.
:param name: Name of event
:param callback: Event callback
:param official_only: Only allow official callbacks
:return: Name of registered callback
"""
def event_callback(*args):
engine.debug('calling callback for event {}'.format(name))
callback(*args)
if official_only and not is_official(name):
raise RuntimeError('tried to register an unofficial event: {name}. try events.event("{name}", official_only=False).'.format(name=name))
callback_name = callbacks.register(event_callback, prefix='event_{}'.format(name))
aggressor.on(name, event_callback)
return callback_name
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)
