diff --git a/assemblyline_client/submit.py b/assemblyline_client/submit.py index d67515b..72ec848 100644 --- a/assemblyline_client/submit.py +++ b/assemblyline_client/submit.py @@ -5,6 +5,7 @@ import select import json import io +from copy import deepcopy from errno import EPIPE from getpass import getpass @@ -18,18 +19,40 @@ from assemblyline_client.v4_client.client import Client as Client4 from assemblyline_client import get_client, __version__ as client_version -from assemblyline_client.v4_client.common.utils import ClientError, get_random_id, get_id_from_path +from assemblyline_client.v4_client.common.utils import ( + ClientError, + get_random_id, + get_id_from_path, +) ASYNC_LOCK = Lock() __version__ = "al_submit v%s" % client_version al_result_to_text = None +DEFAULT_CONFIG = { + "auth": { + "user": None, + "password": None, + "apikey": None, + "insecure": False, + "cert": None, + }, + "server": { + "transport": "https", + "host": "localhost", + "port": 443, + "cert": None, + }, +} +DEFAULT_CONFIG_PATH = "~/.al/submit.cfg" +DEFAULT_TOML_PATH = "~/.al/submit.toml" + # noinspection PyCallingNonCallable def result_to_text(data): if sys.version_info < (3, 0): - results = [x.decode('utf-8') for x in al_result_to_text(data)] + results = [x.decode("utf-8") for x in al_result_to_text(data)] else: results = al_result_to_text(data) return "\n".join(results) @@ -38,7 +61,7 @@ def result_to_text(data): def get_details_from_key(key): file_hash = key[:64] key = key[65:] - name = key[:key.index(".")] + name = key[: key.index(".")] return file_hash, name @@ -59,19 +82,23 @@ def send(client, path, output, options=None, **kw): if options is None: options = {} name = basename(path) - verbose = options.get('verbose', False) + verbose = options.get("verbose", False) try: submission = client.submit(path=path, **kw) - sid = submission.get('sid', None) or submission.get('submission', {}).get('sid', None) + sid = submission.get("sid", None) or submission.get("submission", {}).get( + "sid", None + ) if not sid: - sys.stderr.write("!!ERROR!! Could not find the sid opf the submitted file.\n") + sys.stderr.write( + "!!ERROR!! Could not find the sid opf the submitted file.\n" + ) return False if verbose: sys.stderr.write("File %s submitted for analysis [sid: %s]\n" % (name, sid)) - wq_id = client.live.setup_watch_queue(sid)['wq_id'] + wq_id = client.live.setup_watch_queue(sid)["wq_id"] if verbose: sys.stderr.write("\tListening for incoming results (WQ_ID: %s)\n" % wq_id) @@ -80,7 +107,7 @@ def send(client, path, output, options=None, **kw): while not done: msgs = client.live.get_message_list(wq_id) for m in msgs: - if m['type'] == "start": + if m["type"] == "start": if verbose: sys.stderr.write("\tProcessing...\n") @@ -90,25 +117,30 @@ def send(client, path, output, options=None, **kw): # request to start a watch queue for a file it # hasn't received it yet. Check completion via # submission.is_completed api, continue listening if not completed. - elif m['type'] == "stop" and not start_msg_received: + elif m["type"] == "stop" and not start_msg_received: if client.submission.is_completed(sid): compute_results(client, sid, output, verbose, name, options) done = True break else: - wq_id = client.live.setup_watch_queue(sid)['wq_id'] + wq_id = client.live.setup_watch_queue(sid)["wq_id"] if verbose: - sys.stderr.write("\tSubmission hasn't started on the server yet (new WQ_ID: %s)\n" % wq_id) + sys.stderr.write( + "\tSubmission hasn't started on the server yet (new WQ_ID: %s)\n" + % wq_id + ) - elif m['type'] == "stop": + elif m["type"] == "stop": compute_results(client, sid, output, verbose, name, options) done = True break elif m["type"] == "cachekey" or m["type"] == "cachekeyerr": file_hash, srv_name = get_details_from_key(m["msg"]) if verbose: - m_type = 'ERROR' if m['type'] == 'cachekeyerr' else 'SUCCESS' - sys.stderr.write("\t\t[x] %s (%s) - %s\n" % (srv_name, file_hash, m_type)) + m_type = "ERROR" if m["type"] == "cachekeyerr" else "SUCCESS" + sys.stderr.write( + "\t\t[x] %s (%s) - %s\n" % (srv_name, file_hash, m_type) + ) else: if verbose: sys.stdout.write("%s\n" % m) @@ -121,9 +153,11 @@ def send(client, path, output, options=None, **kw): sys.stderr.write("!!ERROR!! Authentication to the server failed.\n") elif e.status_code == 403: data = json.loads(e) - sys.stderr.write("!!ERROR!! %s\n" % data['api_error_message']) + sys.stderr.write("!!ERROR!! %s\n" % data["api_error_message"]) elif e.status_code == 400 and "File empty" in str(e): - sys.stderr.write("!!ERROR!! Failed to submit '%s' skipped because it is empty.\n" % path) + sys.stderr.write( + "!!ERROR!! Failed to submit '%s' skipped because it is empty.\n" % path + ) else: raise return False @@ -146,80 +180,122 @@ def _main(arguments): signal(SIGINT, SIG_DFL) if sys.platform.startswith("linux"): from signal import SIGPIPE + signal(SIGPIPE, SIG_DFL) - user = None - pw = None - cert = None - apikey = None - transport = "https" - host = "localhost" - port = 443 kw = {} - verify = True - - config = ConfigParser() - config.read([expanduser("~/.al/submit.cfg")]) - for section in config.sections(): - if section == "auth": - if 'user' in config.options('auth') and config.get('auth', 'user'): - user = config.get('auth', 'user') - if 'password' in config.options('auth') and config.get('auth', 'password'): - pw = config.get('auth', 'password') - if 'cert' in config.options('auth') and config.get('auth', 'cert'): - cert = config.get('auth', 'cert') - if 'apikey' in config.options('auth') and config.get('auth', 'apikey'): - apikey = config.get('auth', 'apikey') - if 'insecure' in config.options('auth') and config.get('auth', 'insecure'): - verify = config.get('auth', 'insecure').lower() not in ['true', 'yes'] - elif section == "server": - if 'transport' in config.options('server') and config.get('server', 'transport'): - transport = config.get('server', 'transport') - if 'host' in config.options('server') and config.get('server', 'host'): - host = config.get('server', 'host') - if 'port' in config.options('server') and config.get('server', 'port'): - port = config.get('server', 'port') - if 'cert' in config.options('server') and config.get('server', 'cert'): - verify = config.get('server', 'cert') - else: - sys.stderr.write("The configuration section '[%s]' is invalid. " - "The only valid sections are '[auth]' and '[server]'.\n" % section) + try: + al_config = read_toml() + except (ModuleNotFoundError, FileNotFoundError): + al_config = read_config() - server = "%s://%s:%s" % (transport, host, port) + al_auth = al_config["auth"] + al_server = al_config["server"] # parse the command line args from argparse import ArgumentParser + parser = ArgumentParser(description=description_string) - parser.add_argument('files', metavar='file/dir', nargs='+') - parser.add_argument('-v', '--version', action='version', version=__version__) - parser.add_argument('-q', '--quiet', action='store_true', help='Runs in quiet mode') - parser.add_argument('-a', '--async', dest='async_command', action='store_true', - help='Run in asynchronized mode (uses ingest API).') - parser.add_argument('-n', '--no-output', action='store_true', - help='Only works in conjunction with -a. Ingests the file and does not wait for the output.') - parser.add_argument('-i', '--insecure', action='store_true', default=not verify, - help='Skip server cert validation. DEFAULT: insecure in auth section of ~/.al/submit.cfg') - parser.add_argument('-t', '--text', action='store_true', help='Dumps results as text instead of json.') - parser.add_argument('-d', '--run-dynamic', action='store_true', - help='Adds Dynamic Analysis to the list of service to run.') - parser.add_argument('-u', '--user', default=user, metavar='"user"', - help='username to be used to connect to AL. DEFAULT: user in auth section of ~/.al/submit.cfg') - parser.add_argument('-p', '--password', default=pw, metavar='"MYPASSWORD"', - help='password of the user. DEFAULT: password in auth section of ~/.al/submit.cfg') - parser.add_argument('-o', '--output-file', metavar='"/home/user/output.txt"', - help='File to write the results to. DEFAULT: stdout') - parser.add_argument('-s', '--server', default=server, metavar='"https://localhost:443"', - help='Server to connect to. DEFAULT: transport://host:port in ' - 'server section of ~/.al/submit.cfg') - parser.add_argument('-c', '--cert', default=cert, metavar='"/path/to/pki.pem"', - help='Client cert used to connect to server. DEFAULT: cert in auth section of ~/.al/submit.cfg') - parser.add_argument('-k', '--apikey', default=apikey, metavar='"MY_RANDOM_API_KEY"', - help='apikey to use for the user to login. DEFAULT: apikey in auth section of ~/.al/submit.cfg') - parser.add_argument('-j', '--json-params', metavar='"{ ... }"', help='A JSON dictionary of submission parameters.') - parser.add_argument('-m', '--metadata', metavar='"{ ... }"', help='A JSON dictionary of submission metadata.') - parser.add_argument('--srv-spec', metavar='"{ ... }"', help='A JSON dictionary of service specific parameters.') - parser.add_argument('--server-crt', metavar='"/path/to/server.crt"', - help='DEFAULT: cert in server section of ~/.al/submit.cfg') + parser.add_argument("files", metavar="file/dir", nargs="+") + parser.add_argument("-v", "--version", action="version", version=__version__) + parser.add_argument("-q", "--quiet", action="store_true", help="Runs in quiet mode") + parser.add_argument( + "-a", + "--async", + dest="async_command", + action="store_true", + help="Run in asynchronized mode (uses ingest API).", + ) + parser.add_argument( + "-n", + "--no-output", + action="store_true", + help="Only works in conjunction with -a. Ingests the file and does not wait for the output.", + ) + parser.add_argument( + "-i", + "--insecure", + action="store_true", + default=al_auth["insecure"], + help=f"Skip server cert validation. DEFAULT: insecure in auth section of {DEFAULT_CONFIG_PATH}", + ) + parser.add_argument( + "-t", + "--text", + action="store_true", + help="Dumps results as text instead of json.", + ) + parser.add_argument( + "-d", + "--run-dynamic", + action="store_true", + help="Adds Dynamic Analysis to the list of service to run.", + ) + parser.add_argument( + "-u", + "--user", + default=al_auth["user"], + metavar='"user"', + help=f"username to be used to connect to AL. DEFAULT: user in auth section of {DEFAULT_CONFIG_PATH}", + ) + parser.add_argument( + "-p", + "--password", + default=al_auth["password"], + metavar='"MYPASSWORD"', + help=f"password of the user. DEFAULT: password in auth section of {DEFAULT_CONFIG_PATH}", + ) + parser.add_argument( + "-o", + "--output-file", + metavar='"/home/user/output.txt"', + help="File to write the results to. DEFAULT: stdout", + ) + parser.add_argument( + "-s", + "--server", + default=al_server["url"], + metavar='"https://localhost:443"', + help="Server to connect to. DEFAULT: transport://host:port in " + f"server section of {DEFAULT_CONFIG_PATH}", + ) + parser.add_argument( + "-c", + "--cert", + default=al_auth["cert"], + metavar='"/path/to/pki.pem"', + help=f"Client cert used to connect to server. DEFAULT: cert in auth section of {DEFAULT_CONFIG_PATH}", + ) + parser.add_argument( + "-k", + "--apikey", + default=al_auth["apikey"], + metavar='"MY_RANDOM_API_KEY"', + help=f"apikey to use for the user to login. DEFAULT: apikey in auth section of {DEFAULT_CONFIG_PATH}", + ) + parser.add_argument( + "-j", + "--json-params", + metavar='"{ ... }"', + help="A JSON dictionary of submission parameters.", + ) + parser.add_argument( + "-m", + "--metadata", + metavar='"{ ... }"', + help="A JSON dictionary of submission metadata.", + ) + parser.add_argument( + "--srv-spec", + metavar='"{ ... }"', + help="A JSON dictionary of service specific parameters.", + ) + parser.add_argument( + "--server-crt", + default=al_server["cert"], + metavar='"/path/to/server.crt"', + help=f"DEFAULT: cert in server section of {DEFAULT_CONFIG_PATH}", + ) params = parser.parse_args(arguments) @@ -231,23 +307,22 @@ def _main(arguments): dynamic = params.run_dynamic user = params.user cert = params.cert - pw = params.password + password = params.password apikey = params.apikey - - if params.insecure: - verify = False - else: - if params.server_crt: - verify = params.server_crt + verify = not params.insecure + if not params.insecure and params.server_crt: + verify = params.server_crt if not cert and not user: sys.stderr.write("This server requires authentication...\n") sys.exit(1) - if user and not pw and not apikey: + if user and not password and not apikey: if verbose: - sys.stderr.write("You specified a username without a password. What is your password?\n") - pw = getpass() + sys.stderr.write( + "You specified a username without a password. What is your password?\n" + ) + password = getpass() output = params.output_file @@ -266,11 +341,14 @@ def _main(arguments): server = params.server if not server: - sys.stderr.write("!!ERROR!! No server specified, -s option is mandatory.\n\n%s" % parser.format_help()) + sys.stderr.write( + "!!ERROR!! No server specified, -s option is mandatory.\n\n%s" + % parser.format_help() + ) return -1 if params.metadata: - kw['metadata'] = json.loads(params.metadata) + kw["metadata"] = json.loads(params.metadata) if params.json_params: kw["params"] = json.loads(params.json_params) @@ -283,18 +361,25 @@ def _main(arguments): api_auth = None if user and apikey: api_auth = (user, apikey) - elif user and pw: - auth = (user, pw) + elif user and password: + auth = (user, password) options = { - 'verbose': verbose, - 'json_output': json_output, + "verbose": verbose, + "json_output": json_output, } read_from_pipe = False if sys.platform.startswith("linux") or sys.platform.startswith("freebsd"): try: - if select.select([sys.stdin, ], [], [], 0.0)[0]: + if select.select( + [ + sys.stdin, + ], + [], + [], + 0.0, + )[0]: read_from_pipe = True except io.UnsupportedOperation: # stdin has probably been replaced with a non-file python object @@ -306,32 +391,38 @@ def _main(arguments): return 0 try: - client = get_client(server, apikey=api_auth, auth=auth, cert=cert, verify=verify) + client = get_client( + server, apikey=api_auth, auth=auth, cert=cert, verify=verify + ) if isinstance(client, Client4): - from assemblyline_client.v4_client.common.submit_utils import al_result_to_text + from assemblyline_client.v4_client.common.submit_utils import ( + al_result_to_text, + ) else: from assemblyline_client.v3_client.utils import al_result_to_text except ClientError as e: if e.status_code == 401: sys.stderr.write("!!ERROR!! Authentication to the server failed.\n") elif e.status_code == 495: - sys.stderr.write("!!ERROR!! Invalid SSL connection to the server:\n\t%s\n" % e) + sys.stderr.write( + "!!ERROR!! Invalid SSL connection to the server:\n\t%s\n" % e + ) else: raise return 1 if dynamic: p = client.user.submission_params("__CURRENT__") - if "Dynamic Analysis" not in p['services']['selected']: - p['services']['selected'].append("Dynamic Analysis") + if "Dynamic Analysis" not in p["services"]["selected"]: + p["services"]["selected"].append("Dynamic Analysis") - if 'params' in kw: - p.update(kw['params']) + if "params" in kw: + p.update(kw["params"]) - kw['params'] = p + kw["params"] = p if async_command and not no_output: - kw['nq'] = "al_submit_%s" % get_random_id() + kw["nq"] = "al_submit_%s" % get_random_id() # sanity check path if len(args) == 0 and read_from_pipe: @@ -341,12 +432,12 @@ def _main(arguments): break line = line.strip() - if line == '-': - line = '/dev/stdin' + if line == "-": + line = "/dev/stdin" if async_command: - kw.setdefault('metadata', {}) - kw['metadata']['al_submit_id'] = get_id_from_path(line) + kw.setdefault("metadata", {}) + kw["metadata"]["al_submit_id"] = get_id_from_path(line) send_async(client, line, verbose=verbose, **kw) else: send(client, line, output, options, **kw) @@ -355,8 +446,8 @@ def _main(arguments): file_list = [] for arg in args: - if arg == '-': - file_list.append('/dev/stdin') + if arg == "-": + file_list.append("/dev/stdin") elif not exists(arg): sys.stderr.write("!!ERROR!! %s => File does not exist.\n" % arg) ret_val = 1 @@ -376,13 +467,15 @@ def _main(arguments): for input_file in file_list: if async_command: - kw.setdefault('metadata', {}) - kw['metadata']['al_submit_id'] = get_id_from_path(input_file) + kw.setdefault("metadata", {}) + kw["metadata"]["al_submit_id"] = get_id_from_path(input_file) if not send_async(client, input_file, verbose=verbose, **kw): with ASYNC_LOCK: queued_files.remove(get_id_from_path(input_file)) if verbose: - sys.stderr.write("\tWARNING: Could not send file %s.\n" % input_file) + sys.stderr.write( + "\tWARNING: Could not send file %s.\n" % input_file + ) ret_val = 1 else: if not send(client, input_file, output, options, **kw): @@ -393,17 +486,85 @@ def _main(arguments): if ret_val != 0 and len(file_list) > 1: if verbose: - sys.stderr.write("\n** WARNING: al_submit encountered some " - "errors while processing multiple files. **\n") + sys.stderr.write( + "\n** WARNING: al_submit encountered some " + "errors while processing multiple files. **\n" + ) return ret_val +def read_toml(config_path=DEFAULT_TOML_PATH) -> dict: + """ + Generate default configuration from TOML + """ + + try: + from tomllib import load + except ModuleNotFoundError: + from tomli import load + + al_config = deepcopy(DEFAULT_CONFIG) + + with open(expanduser(config_path), "rb") as tf: + conf_dict = load(tf) + al_config.update(conf_dict) + add_config_url(al_config) + + return al_config + + +def read_config(config_path=DEFAULT_CONFIG_PATH) -> dict: + """ + Generate default configuration + """ + + al_config = deepcopy(DEFAULT_CONFIG) + + config = ConfigParser() + config.read([expanduser(config_path)]) + for section in config.sections(): + if section == "auth": + for key in ("user", "password", "cert", "apikey", "insecure"): + if key in config.options(section) and config.get(section, key): + if key == "insecure": + al_config[section][key] = config.get( + section, key + ).lower().strip("\"'") in ["true", "yes"] + continue + al_config[section][key] = config.get(section, key).strip("\"'") + elif section == "server": + for key in ("transport", "host", "port", "cert"): + if key in config.options(section) and config.get(section, key): + al_config[section][key] = config.get(section, key).strip("\"'") + elif section == "incident": + pass + else: + sys.stderr.write( + "The configuration section '[%s]' is invalid. " + "The only valid sections are '[auth]' and '[server]'.\n" + "The section '[incident]' is reserved for assemblyline-incident-manager.\n" + % section + ) + + add_config_url(al_config["server"]) + + return al_config + + +def add_config_url(al_server): + al_server["url"] = "%s://%s:%s" % ( + al_server.get("transport", "https"), + al_server.get("host", "localhost"), + al_server.get("port", 443), + ) + + def send_async(client, path, verbose=False, **kw): try: if verbose: sys.stderr.write("Sending file %s for analysis...\n" % path) - client.ingest(path=path, ingest_type='AL_SUBMIT', **kw) + client.ingest(path=path, ingest_type="AL_SUBMIT", **kw) return True except ClientError: return False @@ -411,17 +572,15 @@ def send_async(client, path, verbose=False, **kw): def start_result_thread(client, queued_files, output, options, **kw): output_thread = Thread( - target=result_thread, - args=(client, queued_files, output, options), - kwargs=kw + target=result_thread, args=(client, queued_files, output, options), kwargs=kw ) output_thread.start() return output_thread def result_thread(client, queued_files, output, options, **kw): - nq = kw['nq'] - verbose = options.get('verbose', False) + nq = kw["nq"] + verbose = options.get("verbose", False) while len(queued_files) != 0: if verbose: @@ -429,20 +588,26 @@ def result_thread(client, queued_files, output, options, **kw): msgs = client.ingest.get_message_list(nq) for msg in msgs: - sid = msg.get('submission', {}).get('sid', None) or msg.get('alert', {}).get('sid', None) + sid = msg.get("submission", {}).get("sid", None) or msg.get( + "alert", {} + ).get("sid", None) if not sid: - sys.stderr.write("!!ERROR!! Could not find the sid of the submitted " - "file in the message.\n{}".format(msg)) + sys.stderr.write( + "!!ERROR!! Could not find the sid of the submitted " + "file in the message.\n{}".format(msg) + ) continue try: # v4 structure - cur_file = msg['submission']['files'][0]['name'] - submission_id = msg['submission']['metadata']['al_submit_id'] + cur_file = msg["submission"]["files"][0]["name"] + submission_id = msg["submission"]["metadata"]["al_submit_id"] except KeyError: # v3 structure - cur_file = msg.get('metadata', {}).get('filename', None) or msg['sha256'] - submission_id = msg['metadata']['al_submit_id'] + cur_file = ( + msg.get("metadata", {}).get("filename", None) or msg["sha256"] + ) + submission_id = msg["metadata"]["al_submit_id"] with ASYNC_LOCK: try: @@ -451,7 +616,10 @@ def result_thread(client, queued_files, output, options, **kw): pass if verbose: - sys.stderr.write("\tFile '%s' complete. Fetching results for submission ID: %s...\n" % (cur_file, sid)) + sys.stderr.write( + "\tFile '%s' complete. Fetching results for submission ID: %s...\n" + % (cur_file, sid) + ) final_results = client.submission.full(sid) if output: @@ -466,12 +634,21 @@ def result_thread(client, queued_files, output, options, **kw): def write_file(data, path, infile, verbose=False, json_output=True): with open(path, "ab") as out_file: if json_output: - - out_file.write("[{}] {} <==> {}\n".format(datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'), - infile, json.dumps(data, separators=(',', ':'))).encode()) + out_file.write( + "[{}] {} <==> {}\n".format( + datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + infile, + json.dumps(data, separators=(",", ":")), + ).encode() + ) else: - out_file.write("[{}] {}\n\n{}\n\n--------\n\n".format(datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'), - infile, result_to_text(data)).encode()) + out_file.write( + "[{}] {}\n\n{}\n\n--------\n\n".format( + datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + infile, + result_to_text(data), + ).encode() + ) if verbose: sys.stderr.write("%s => Resulting file saved to %s\n" % (infile, path))