def restart_tasks(tasks: List[Task]) -> None:
identity = "karton.dashboard-retry"
producer = Producer(identity=identity)
for task in tasks:
# spawn a new task and mark the original one as finished
producer.send_task(task.fork_task())
karton.backend.set_task_status(task=task,
status=TaskState.FINISHED,
consumer=identity)
def upload():
producer = Producer(conf)
with NamedTemporaryFile() as f:
request.files["file"].save(f.name)
with open(f.name, "rb") as fr:
sample = Resource("sample", fr.read())
task = Task({"type": "sample", "stage": "recognized", "platform": "win32"})
task.add_payload("override_uid", task.uid)
# Add analysis timeout to task
timeout = request.form.get("timeout")
if timeout:
task.add_payload("timeout", int(timeout))
# Add filename override to task
if request.form.get("file_name"):
filename = request.form.get("file_name")
else:
filename = request.files["file"].filename
if not re.fullmatch(
r"^((?![\\/><|:&])[\x20-\xfe])+\.(?:dll|exe|ps1|bat|doc|docm|docx|dotm|xls|xlsx|xlsm|xltx|xltm|ppt|pptx|vbs|js|jse|hta|html|htm)$",
filename,
flags=re.IGNORECASE,
):
return jsonify({"error": "invalid file_name"}), 400
task.add_payload("file_name", os.path.splitext(filename)[0])
# Extract and add extension
extension = os.path.splitext(filename)[1][1:]
if extension:
task.headers["extension"] = extension
# Add startup command to task
start_command = request.form.get("start_command")
if start_command:
task.add_payload("start_command", start_command)
# Add plugins to task
plugins = request.form.get("plugins")
if plugins:
plugins = json.loads(plugins)
task.add_payload("plugins", plugins)
task.add_resource("sample", sample)
producer.send_task(task)
return jsonify({"task_uid": task.uid})
def submit_main(cls):
parser = cls.args_parser()
args = parser.parse_args()
conf_path = os.path.join(ETC_DIR, "config.ini")
config = patch_config(Config(conf_path))
with open(args.tests) as tests:
testcases = [TestCase(**case) for case in json.load(tests)]
root_uids = []
for test in testcases:
sample = test.get_sample()
sys.stderr.write(f"Submitting {test.sha256}\n")
t = Task(headers=dict(type="sample-test", platform="win64"))
t.add_payload("sample", Resource("malwar", sample))
t.add_payload("testcase", test.to_json())
if args.timeout:
t.add_payload("timeout", args.timeout)
p = Producer(config)
p.send_task(t)
root_uids.append(t.root_uid)
consumer = RegressionTester(config)
results = {}
with tqdm(total=len(root_uids)) as pbar:
while len(results) != len(root_uids):
for root_uid in cls.get_finished_tasks(consumer.backend,
root_uids):
if root_uid not in results:
res = json.load(
consumer.backend.minio.get_object(
"draktestd", root_uid))
results[root_uid] = res
print(json.dumps(results[root_uid]))
pbar.update(1)
time.sleep(1)
print(json.dumps(list(results.values())))
def main():
parser = argparse.ArgumentParser(description="Push sample to the karton")
parser.add_argument("sample", help="Path to the sample")
parser.add_argument(
"--start_command",
help="e.g. start %f, %f will be replaced by file name",
required=False,
)
parser.add_argument(
"--timeout",
default=600,
type=int,
help="analysis timeout in seconds",
required=False,
)
args = parser.parse_args()
conf = patch_config(Config(os.path.join(ETC_DIR, "config.ini")))
producer = Producer(conf)
task = Task({"type": "sample", "stage": "recognized", "platform": "win32"})
with open(args.sample, "rb") as f:
sample = Resource("sample", f.read())
task.add_resource("sample", sample)
# Add filename
filename = os.path.basename(args.sample)
task.add_payload("file_name", os.path.splitext(filename)[0])
# Extract and add extension
extension = os.path.splitext(filename)[1][1:]
if extension:
task.headers["extension"] = extension
if args.start_command is not None:
task.add_payload("start_command", args.start_command)
if args.timeout is not None:
task.add_payload("timeout", args.timeout)
producer.send_task(task)
def send_config_to_karton(cfg: Config) -> str:
producer = Producer(identity="karton.mwdb",
config=KartonConfig(config.karton.config_path))
task = Task(headers={
"type": "config",
"kind": cfg.config_type,
"family": cfg.family
},
payload={
"config":
cfg.cfg,
"dhash":
cfg.dhash,
"attributes":
cfg.get_metakeys(as_dict=True, check_permissions=False)
})
producer.send_task(task)
cfg.add_metakey("karton", task.root_uid, check_permissions=False)
logger.info("Configuration sent to karton with %s", task.root_uid)
return task.root_uid
def send_file_to_karton(file: File) -> str:
try:
path = file.get_path()
tmpfile = None
except Exception:
# If get_path doesn't work: download content to NamedTemporaryFile
# It won't work if we use S3 storage and try to reanalyze
# existing file (not uploaded within the same request).
tmpfile = tempfile.NamedTemporaryFile()
file_stream = file.open()
shutil.copyfileobj(file_stream, tmpfile)
File.close(file_stream)
path = tmpfile.name
producer = Producer(identity="karton.mwdb",
config=KartonConfig(config.karton.config_path))
feed_quality = g.auth_user.feed_quality
task_priority = TaskPriority.NORMAL if feed_quality == "high" else TaskPriority.LOW
task = Task(headers={
"type": "sample",
"kind": "raw",
"quality": feed_quality
},
payload={
"sample":
Resource(file.file_name, path=path, sha256=file.sha256),
"attributes":
file.get_metakeys(as_dict=True, check_permissions=False)
},
priority=task_priority)
producer.send_task(task)
if tmpfile is not None:
tmpfile.close()
file.add_metakey("karton", task.root_uid, check_permissions=False)
logger.info("File sent to karton with %s", task.root_uid)
return task.root_uid
def get_karton_producer() -> Producer:
return Producer(identity="karton.mwdb",
config=KartonConfig(app_config.karton.config_path))