async def _get_documents(self, sha, sca, cmd, args, fmt, otpt):
documents = []
cur = db.async_command_collection.select_many(sha256_digest=sha,
scale=sca,
command=cmd,
args=args,
sort="timestamp")
while await cur.fetch_next:
doc = cur.next_object()
doc = schema.CommandSchema().load(doc)
try: # Ignore output for missing scales and/or commands
scale = scale_manager.get_scale(doc['scale'])
commands = scale_manager.get_component(
scale, enums.ScaleComponent.COMMANDS)
except Exception as err:
print("%s - %s" % (doc['scale'], err)) # TODO: Output to log
continue
output = None
if '_output_id' in doc and doc['_output_id']:
output = await db.async_command_output_collection.get(
doc['_output_id'])
doc = schema.CommandSchema().dump(doc)
try:
if otpt:
doc['output'] = commands.snake.format(fmt, cmd, output)
doc['format'] = fmt
except (SnakeError, TypeError) as err:
print("%s - %s" % (doc['scale'], err)) # TODO: Output to log
continue
documents += [doc]
return documents
async def post(self, scale):
if not self.request.body:
self.write_warning("scale/interface - no request body found", 422,
scale)
self.finish()
return
data = escape.json_decode(self.request.body)
try:
data = self.InterfaceSchema().dump(
self.InterfaceSchema().load(data))
except exceptions.ValidationError as err:
self.write_warning(
self.json_decode(('%s' % err.messages).replace("'", '"')), 422)
self.finish()
return
document = await db.async_file_collection.select(data['sha256_digest'])
if not document:
self.write_warning("scale/interface - no sample for given data",
404, data)
self.finish()
return
# Get the push/pull and args
_scale = scale_manager.get_scale(scale)
interface = scale_manager.get_component(_scale,
enums.ScaleComponent.INTERFACE)
command = scale_manager.get_interface_command(interface, data['type'],
data['command'])
data['timestamp'] = datetime.utcnow()
# Execute command
# TODO: Handle status don't always chuck errors...
try:
loop = asyncio.get_event_loop()
output = await loop.run_in_executor(None, command, data['args'],
data['sha256_digest'])
# output = p(data['args'], data['sha256_digest'])
except exceptions.ValidationError as err:
self.write_warning(
self.json_decode(
('{"args": %s}' % err.messages).replace("'", '"')), 422)
self.finish()
return
except error.SnakeError as err:
self.write_warning("%s" % err, err.status_code, data)
self.finish()
return
# Run formating
data['output'] = interface.snake.format(data['format'],
data['command'], output)
self.jsonify({"interface": data})
self.finish()
async def post(self, data):
# Check that there is a file for this hash
document = await db.async_file_collection.select(data['sha256_digest'])
if not document:
self.write_warning("no sample for given data", 404, data)
self.finish()
return
# Check scale support
try:
scale = scale_manager.get_scale(data['scale'],
document['file_type'])
commands = scale_manager.get_component(
scale, enums.ScaleComponent.COMMANDS)
cmd = commands.snake.command(data['command'])
except SnakeError as err:
self.write_warning("%s" % err, 404, data)
self.finish()
return
# Validate arguments as to not waste users time, yes this is also done on execution
result, args = validate_args(cmd, data['args'])
if not result:
self.write_warning(args, 422, data)
self.finish()
return
data['args'] = args
# Queue command
try:
document = await route_support.queue_command(data)
except SnakeError as err:
self.write_warning("%s" % err, 500, data)
self.finish()
return
document = schema.CommandSchema().load(document)
output = None
if document['_output_id']:
output = await db.async_command_output_collection.get(
document['_output_id'])
try:
document['output'] = commands.snake.format(data['format'],
document['command'],
output)
document['format'] = data['format']
except SnakeError as err:
self.write_warning("%s" % err, 404, data)
self.finish()
return
# Dump and finish
document = schema.CommandSchema().dump(document)
self.jsonify({"command": document})
self.finish()
async def get(self, scale):
try:
_scale = scale_manager.get_scale(scale)
commands = scale_manager.get_component(
_scale, enums.ScaleComponent.COMMANDS)
except error.SnakeError as err:
self.write_warning("scale - %s" % err, err.status_code, scale)
self.finish()
return
self.jsonify({"commands": commands.snake.info()})
self.finish()
async def get(self, scale):
reload = self.get_argument('reload', "false")
if reload.lower() == "true": # Ignore all other values
scale_manager.reload_scales()
try:
_scale = scale_manager.get_scale(scale)
except error.SnakeError as err:
self.write_warning("scale - %s" % err, err.status_code, scale)
self.finish()
return
self.jsonify({"scale": _scale.info()})
self.finish()
async def get(self, data):
# NOTE: Tornado/Marshmallow does not like Dict in args, will have to parse manually
# TODO: Use marshmallow validation
if 'args' in self.request.arguments and self.request.arguments['args']:
data['args'] = json.loads(self.request.arguments['args'][0])
else:
data['args'] = {}
document = await db.async_command_collection.select(
data['sha256_digest'], data['scale'], data['command'],
data['args'])
if not document:
self.write_warning("no output for given data", 404, data)
self.finish()
return
if document['status'] == enums.Status.ERROR:
self.write_warning("%s" % document['output'], 404, data)
self.finish()
return
document = schema.CommandSchema().load(document)
output = None
if document['_output_id']:
output = await db.async_command_output_collection.get(
document['_output_id'])
try:
scale = scale_manager.get_scale(data['scale'])
commands = scale_manager.get_component(
scale, enums.ScaleComponent.COMMANDS)
if data['output']:
document['output'] = commands.snake.format(
data['format'], document['command'], output)
document['format'] = data['format']
except (SnakeError, TypeError) as err:
self.write_warning("%s" % err, 404, data)
self.finish()
return
document = schema.CommandSchema().dump(document)
self.jsonify({'command': document})
self.finish()
async def get(self, scale):
_scale = scale_manager.get_scale(scale)
interface = scale_manager.get_component(_scale,
enums.ScaleComponent.INTERFACE)
self.jsonify({"interface": interface.snake.info()})
self.finish()
async def post(self, scale): # pylint: disable=too-many-locals
if not self.request.body:
self.write_warning("scale/upload - no request body found", 422,
scale)
self.finish()
return
data = escape.json_decode(self.request.body)
# Validate args
try:
data = self.UploadSchema().dump(self.UploadSchema().load(data))
except exceptions.ValidationError as err:
self.write_warning(
self.json_decode(('%s' % err.messages).replace("'", '"')), 422)
self.finish()
return
scale_ = scale_manager.get_scale(scale)
upload = scale_manager.get_component(scale_,
enums.ScaleComponent.UPLOAD)
# Validate arguments and update
upld_args = upload.arguments()
try:
if upld_args:
data['args'] = schema.Schema(fields=upld_args).load(
data['args'])
except exceptions.ValidationError as err:
self.write_warning(
self.json_decode(
('{"args": %s}' % err.messages).replace("'", '"')), 422)
self.finish()
return
# Get the file
with tempfile.TemporaryDirectory(dir=path.abspath(
path.expanduser(snake_config['cache_dir']))) as temp_dir:
loop = asyncio.get_event_loop()
f_name = await loop.run_in_executor(None, upload.upload,
data['args'], temp_dir)
f_path = path.join(temp_dir, f_name)
# Extract if required, zip only
if data['extract']:
f_path = await route_support.unzip_file(
f_path, data['password'])
f_name = path.basename(f_path)
# Update name if not overriden
if not data['name']:
data['name'] = f_name
# Set submission type
data['submission_type'] = 'upload:{}'.format(scale)
# Check that the file is not empty
if path.getsize(f_path) == 0:
self.write_warning("scale/upload - sample is empty", 422)
self.finish()
return
# Hash the file
sha2 = hashlib.sha256()
with open(f_path, 'rb') as f:
chunk = f.read(4096)
while chunk:
sha2.update(chunk)
chunk = f.read(4096)
sha256_digest = sha2.hexdigest()
# Check if the file already exists
document = await db.async_file_collection.select(sha256_digest)
if document:
document = schema.FileSchema().dump(
schema.FileSchema().load(document))
self.write_warning(
"scale/upload - sample already exists for given sha256 digest",
409, {'sample': document})
self.finish()
return
# Save the file and add it to the database
document = await route_support.store_file(sha256_digest, f_path,
data['file_type'], data)
document = schema.FileSchema().dump(
schema.FileSchema().load(document))
self.jsonify({'sample': document})
self.finish()
async def get(self, scale):
scale_ = scale_manager.get_scale(scale)
upload = scale_manager.get_component(scale_,
enums.ScaleComponent.UPLOAD)
self.jsonify({"upload": upload.snake.info()})
self.finish()
async def post(self, data): # pylint: disable=too-many-locals, too-many-branches, too-many-statements
# XXX: Needs a major clean/rework
if not data:
self.write_warning("commands - no request body found", 422, data)
self.finish()
return
# Find the commands and validate their arguments
for d in data:
# Find the command
try:
s = scale_manager.get_scale(d['scale'])
c = scale_manager.get_component(s,
enums.ScaleComponent.COMMANDS)
cmd = c.snake.command(d['command'])
except ScaleError as err:
self.write_warning(err.message, err.status_code)
self.finish()
return
result, args = validate_args(cmd, d['args'])
if not result:
self.write_warning(self.json_decode(args.replace("'", '"')),
422, data)
self.finish()
return
d['args'] = args
# Validate hashes and validate them against scales
missing = []
unsupported = []
for d in data:
s = scale_manager.get_scale(d['scale'])
for sha in d['sha256_digests']:
if sha.lower() == 'all':
if not s.supports and not len(s.supports) == len(
[x for x in enums.FileType]):
unsupported += [d]
break
elif sha.lower()[:4] == 'all:':
file_type = sha.lower().split(':')[1]
if file_type == 'file':
ft = enums.FileType.FILE
elif file_type == 'memory':
ft = enums.FileType.MEMORY
else:
ft = None
if ft is None or (s.supports and ft not in s.supports):
unsupported += [(sha, s.name)]
break
else:
document = await db.async_file_collection.select(sha)
if not document:
missing += [d]
elif s.supports and document[
'file_type'] not in s.supports: # Check scale support
unsupported += [d]
if missing:
self.write_warning("commands - no sample(s) for given data", 404,
missing)
self.finish()
return
if unsupported:
self.write_warning("commands - command unsupported for given data",
422, unsupported)
self.finish()
return
# Queue commands
documents = []
for d in data:
cmd_dict = {}
for k, v in d.items():
if k != 'sha256_digests':
cmd_dict[k] = v
cmd_dict['asynchronous'] = True
for sha in d['sha256_digests']:
if sha.lower() == 'all':
cursor = db.async_file_collection.select_all()
while await cursor.fetch_next:
cmd_dict['sha256_digest'] = cursor.next_object(
)['sha256_digest']
cmd_d = schema.CommandSchema().load(cmd_dict)
documents += [await route_support.queue_command(cmd_d)]
break
elif sha.lower()[:4] == 'all:':
ft = sha.lower().split(':')[1]
if ft == 'file':
ft = enums.FileType.FILE
elif ft == 'memory':
ft = enums.FileType.MEMORY
cursor = db.async_file_collection.select_many(file_type=ft)
while await cursor.fetch_next:
cmd_dict['sha256_digest'] = cursor.next_object(
)['sha256_digest']
cmd_d = schema.CommandSchema().load(cmd_dict)
documents += [await route_support.queue_command(cmd_d)]
break
else:
cmd_dict['sha256_digest'] = sha
cmd_d = schema.CommandSchema().load(cmd_dict)
documents += [await route_support.queue_command(cmd_d)]
# Dump and finish
documents = schema.CommandSchema(many=True).load(documents)
documents = schema.CommandSchema(many=True).dump(documents)
self.jsonify({"commands": documents})
self.finish()