| |
| |
| 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!") |
|
|
