|
|
|
|
|
import torch |
|
import pickle_inspector |
|
import sys |
|
from pathlib import Path |
|
|
|
debug = len(sys.argv) == 3 |
|
|
|
dir = sys.argv[1] |
|
print("checking dir: " + dir) |
|
|
|
BASE_DIR = Path(dir) |
|
EXTENSIONS = {'.pt', '.bin', '.ckpt'} |
|
BAD_CALLS = {'os', 'shutil', 'sys', 'requests', 'net'} |
|
BAD_SIGNAL = {'rm ', 'cat ', 'nc ', '/bin/sh '} |
|
|
|
for path in BASE_DIR.glob(r'**/*'): |
|
if path.suffix in EXTENSIONS: |
|
print("") |
|
print("..." + path.as_posix()) |
|
result = torch.load(path.as_posix(), pickle_module=pickle_inspector.pickle) |
|
result_total = 0 |
|
result_other = 0 |
|
result_calls = {} |
|
result_signals = {} |
|
result_output = "" |
|
|
|
for call in BAD_CALLS: |
|
result_calls[call] = 0 |
|
|
|
for signal in BAD_SIGNAL: |
|
result_signals[signal] = 0 |
|
|
|
for c in result.calls: |
|
for call in BAD_CALLS: |
|
if (c.find(call + ".") == 0): |
|
result_calls[call] += 1 |
|
result_total += 1 |
|
result_output += "\n--- found lib call (" + call + ") ---\n" |
|
result_output += c |
|
result_output += "\n---------------\n" |
|
break |
|
for signal in BAD_SIGNAL: |
|
if (c.find(signal) > -1): |
|
result_signals[signal] += 1 |
|
result_total += 1 |
|
result_output += "\n--- found malicious signal (" + signal + ") ---\n" |
|
result_output += c |
|
result_output += "\n---------------\n" |
|
break |
|
|
|
if ( |
|
c.find("numpy.") != 0 and |
|
c.find("_codecs.") != 0 and |
|
c.find("collections.") != 0 and |
|
c.find("torch.") != 0): |
|
result_total += 1 |
|
result_other += 1 |
|
result_output += "\n--- found non-standard lib call ---\n" |
|
result_output += c |
|
result_output += "\n---------------\n" |
|
|
|
if (result_total > 0): |
|
for call in BAD_CALLS: |
|
print("library call (" + call + ".): " + str(result_calls[call])) |
|
for signal in BAD_SIGNAL: |
|
print("malicious signal (" + signal + "): " + str(result_signals[signal])) |
|
print("non-standard calls: " + str(result_other)) |
|
print("total: " + str(result_total)) |
|
print("") |
|
print("SCAN FAILED") |
|
|
|
if (debug): |
|
print(result_output) |
|
else: |
|
print("SCAN PASSED!") |
|
|
