modescan.nasl

modescan.nasl

include("compat.inc");

if(description)
{
  script_id(999001);
  script_version("1.0");
  script_set_attribute(attribute:"plugin_modification_date", value:"2024/06/25");

  script_name(english: "Execute ModelScan Checks against AI/LLM files");
  script_summary(english: "This opensource nasl plugin downloads & executes the open-source ModelScan 
  tooling to perform tests against the model that exists on the remote target");

  script_set_attribute(attribute:"synopsis", value: "Performs ModelScan checks against model files on the remote target.");
  script_set_attribute(attribute:"description", value: "This opensource nasl plugin downloads & executes the open-source ModelScan 
  tooling to perform tests against the model that exists on the remote target.");
  script_set_attribute(attribute:"solution", value:"Investigate the usage of any model flagged.");
  script_set_attribute(attribute:"see_also", value:"https://github.com/protectai/modelscan");
  script_set_attribute(attribute:"see_also", value:"https://github.com/trailofbits/ml-file-formats?tab=readme-ov-file");
  script_set_attribute(attribute:"risk_factor", value: "None");
  script_set_attribute(attribute:"plugin_publication_date", value: "2024/06/25");
  script_set_attribute(attribute:"plugin_type", value:"local");
  script_end_attributes();

  script_category(ACT_GATHER_INFO);

  script_copyright(english:"This script is Copyright (C) 2024 and is owned by Tom Cappetta, an OpenSource Developer.");
  script_family(english: "General");

  script_dependencies("ssh_get_info.nasl", "os_fingerprint.nasl", "command_builder_init.nasl");
  script_require_keys("Host/local_checks_enabled");
  exit(0);
}

include('ssh_func.inc');
include('telnet_func.inc');
include('hostlevel_funcs.inc');
include('http.inc');
include('local_detection_nix.inc');
include('json.inc');
include("sh_commands_find.inc");
include('command_builder.inc');

var mfiles = make_list(
    '*.pt', '*.pth', '*.bin', '*.mar', '*.pte', '*.pt2', '*.ptl', '*.safetensors', '*.onnx', '*.keras', '*.pb', '*.ckpt', '*.tflite', '*.tfrecords', '*.npy', '*.npz', '*.gguf', '*.ggml', '*.ggmf', '*.ggjt', '*.nc', '*.mleap', '*.coreml', '*.surml', '*.llamafile', '*.prompt', '*.pkl', '*.caffemodel', '*.prototxt', '*.h5', '*.dlc'
    );

function findmodelfiles()
{    
    enable_ssh_wrappers();
    dbg::log(msg: 'searching for models');
    if (!get_kb_item("Host/local_checks_enabled"))
    audit(AUDIT_LOCAL_CHECKS_NOT_ENABLED);

    if ("Linux" >!< get_kb_item_or_exit("Host/uname"))
    audit(AUDIT_OS_NOT, "Linux");

    if (islocalhost())
    {
    if (!defined_func("pread")) audit(AUDIT_FN_UNDEF,"pread");
    info_t = INFO_LOCAL;
    }
    else
    {
    sock_g = ssh_open_connection();
    if (!sock_g) audit(AUDIT_FN_FAIL, 'ssh_open_connection');
    info_t = INFO_SSH;
    }

    var find_args = make_list('/');
    if (thorough_tests)
    {
    find_args = make_list(find_args, '/');
    }
    
    var ext = branch(mfiles);
    find_args = make_list(find_args,'-maxdepth', '9', '-xautofs', '-tenb_fstype_exclusions', '-tenb_path_exclusions', '-type', 'l', '-name', ext);

    var find_output = sh_commands::find(args:find_args, timeout:60);
    #dbg::log(msg:'find_output -> ' + obj_rep(find_output));

    if (find_output[0] == sh_commands::CMD_OK)
    {
    find_output = find_output[1];
    #dbg::log(msg:'CMD_OK -> ' ); 
    if (info_t == INFO_SSH) ssh_close_connection();
    if (!empty_or_null(find_output))
        return find_output;
    else return NULL;
    }
    else if (find_output[0] == sh_commands::CMD_TIMEOUT)
    {
    if (info_t == INFO_SSH) ssh_close_connection();
    exit(1, 'Find command timed out.');
    }
    else
    {
    if (info_t == INFO_SSH) ssh_close_connection();
    #dbg::log(msg:'ELSE output -> ' + obj_rep(find_output));
    exit(1, find_output[1]);
    }
}

function create_venv()
{
    dbg::log(msg:'creating python virtual env in: ' + tmpdir);
    #var pyvenv_result = ldnix::run_cmd_template_wrapper(template:"python3 -m venv " + tmpdir + "; . "+ tmpdir + "/bin/activate && pip install modelscan 'modelscan[ tensorflow, h5py]'");
    var pyvenv_result = ldnix::run_cmd_template_wrapper(template:"rm -rf /tmp/venv_*; python3 -m venv " + tmpdir + "; . "+ tmpdir + "/bin/activate && pip install modelscan");
}

function scanfile(file)
{
    var result = ldnix::run_cmd_template_wrapper(template:". "+ tmpdir + "/bin/activate && modelscan -r json -p " + file);
    var lines = split(result, sep:'\n', keep:FALSE);
    dbg::log(msg:'modelscan result \n[' + file + ']\n' + obj_rep(lines));
    var json;
    var skip2=0;
    foreach var line (lines)
    {
        if (skip2 < 3){skip2 +=1; dbg::log(msg:'moving to line: ' + skip2); continue;}
        #dbg::log(msg:'line --\n' + obj_rep(line));
        json = json + line;
    }
    dbg::log(msg:'json\n||\n' + obj_rep(json) + '\n||\n');
    var jdata = json_read(json);
    dbg::log(msg:'modelscan json result --\n' + obj_rep(jdata));

    var abspath = jdata[0]['summary']['absolute_path'];
    var severity = SECURITY_WARNING;
    var output = jdata[0]['issues'];
    var request = jdata[0]['summary']['absolute_path'];
    var issues = jdata[0]['summary']['total_issues'];
    
    # function security_report_v4(port, proto, severity, extra, cmd, file, request, output, line_limit, rep_extra, attach_type, generic, sqli, xss, xsrf)
    var report = 'Nessus downloaded and ran the opensource AI/LLM modelscan tool.  Here are the results: \n';
    report += strcat('Total issues found: ', issues , '\n');
    report += strcat('ModelFile: ',request , '\n');
    
    foreach var is (output)
    {
        foreach var metadata (keys(is))
        {
            report += metadata + ': \t' +  is[metadata] + '\n';
        }
    }

    var bad = ver_compare(ver:issues, fix:0 );
    dbg::log(msg:'issues found -> ' + issues + '\n bad -> ' + bad);


    if(ver_compare(ver:issues, fix:0 ) > 0)
    #security_report_v4(port:0, severity:SECURITY_HOLE, generic:TRUE, request:obj_rep(request), output:obj_rep(output));
        security_report_v4(port:0, severity:SECURITY_HOLE, extra:report);
    else audit(AUDIT_PACKAGE_NOT_AFFECTED, abspath );
}


get_kb_item_or_exit('Host/local_checks_enabled');

set_kb_item(name:'global_settings/enable_plugin_debugging', value:TRUE);
set_kb_item(name:'global_settings/debug_level', value:'4');
enable_ssh_wrappers();
info_connect();

var tmpdir = '/tmp/venv_' + rand();
create_venv();

dbg::log(msg:'testing modelsearch');
var found_files = findmodelfiles();
if (empty_or_null(found_files))
    audit(AUDIT_PACKAGE_NOT_INSTALLED, 'AI / LLM model files');

var ppath = ldnix::get_command_path(command:'pip');
dbg::log(msg:'pip -> ' + obj_rep(ppath));

var pypath = ldnix::get_command_path(command:'python3');
dbg::log(msg:'python3 -> ' + obj_rep(pypath));

var mpath = ldnix::get_command_path(command:'modelscan');
dbg::log(msg:'modelscan -> ' + obj_rep(mpath));

var mvp = ldnix::file_exists(file:tmpdir + "/bin/modelscan");
dbg::log(msg:'mvp -> ' + obj_rep(mvp));
#exit(0);

if (empty_or_null(ppath) || empty_or_null(pypath))
    audit(AUDIT_PACKAGE_NOT_INSTALLED, 'python / pip');

if (!empty_or_null(found_files)){
    var files = split(found_files, sep:'\n', keep:FALSE);
    dbg::log(msg:'scanning file -> ' + obj_rep(files));
    var file = branch(files);
    scanfile(file:file);
}