TensorAbuse/main.py at master · KramerZJ/TensorAbuse · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
from PersistExt.py_func_op_extract import PyFuncOpExtract
from PersistExt.py_C_analysis import py_C_analysis
import argparse
import re
import os
import sys
from TensorDetect import model, scan
pattern = r'^(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)' \
              r'(-[0-9A-Za-z-]+(\.[0-9A-Za-z-]+)*)?' \
              r'(\+[0-9A-Za-z-]+(\.[0-9A-Za-z-]+)*)?$'

def main(args):
    if args.path and args.version and args.target:
        path = args.path
        if not os.path.exists(path):
            print("Invalid path")
            return

        if re.match(pattern, args.version) is None:
            print("Invalid version format")
            return
        else:
            version = 2 if args.version.split(".")[0]=='2' else 1
        py_func_op_extract = PyFuncOpExtract(version, path)
        py_func_op_extract.analyze_tensorflow_source()
        py_func_op_extract.get_results(args.target)

    elif args.base_path:
        path = args.base_path
        py_C_analysis(path)

    elif args.model:
        path = args.model
        if not os.path.exists(path):
            print("Invalid path!")
            return

        mod = model.Model(path)
        if mod.model_type == model.ModelType.TF_H5:
            sc = scan.H5Scan(mod)
        elif mod.model_type == model.ModelType.TF_SM:
            sc = scan.SavedModelScan(mod)
        sc.scan()
        sc.print_issues()

    else:
        print("Invalid arguments, please provide path, version and target file path!")
        return


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Analyze Tensorflow API or detect malicious Tensorflow models(i.e., h5 or saved_model).")

    group1 = parser.add_argument_group("Tensorflow API analysis group")
    group1.add_argument("-p", "--path", help="Tensorflow source code path")
    group1.add_argument("-v", "--version", help="Tensorflow version")
    group1.add_argument("-t", "--target", help="Target file path")

    group2 = parser.add_argument_group("Malicious model detection group")
    # group2.add_argument("-d", "--detect", help="Model detection", default=1)
    group2.add_argument("-m", "--model", help="Tensorflow model path (i.e., h5 or saved_model)")

    group3 = parser.add_argument_group("PersistExt analyze python-C++ cross language call chain")
    group3.add_argument("-b", "--base_path", help="Path to Tensorflow source code folder")

    args = parser.parse_args()

    # Ensure group1 or group2 parameters are provided as a full set
    if (args.path or args.version or args.target) and not (args.path and args.version and args.target):
        print("Error: Arguments -p, -v, -t must be used together.")
        sys.exit(1)

    elif not (args.path and args.version and args.target) and not (args.model or args.base_path):
        print("Error: Arguments -m or -b must be used.")
        sys.exit(1)

    main(args)