ceo/ceo-bin at master · trailofbits/ceo · 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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
#!/usr/bin/env python2

from argparse import ArgumentParser
from sys import argv, exit, maxint
from random import seed, choice

from ceo.init     import init
from ceo.testcase import Testcase, load_targets, load_testcases
from ceo.exploration import explore
from ceo.test     import test, stats
from ceo.actions  import MinimizeTestcaseAFL, MinimizeInputsAFL
from ceo.parameters  import init_parameter_generators
from ceo.storage   import Storage
from ceo.labels   import lbls
from ceo.reduce  import reduce_inputs, reduce_testcase
from ceo.extraction  import get_features
from ceo.features import features_list

'''

CEO

'''

def parse_arguments():
    ###########################################################################
    # parse arguments
    parser = ArgumentParser(description='')

    parser.add_argument('--seed', type=int, default=None,
                        help="Random seed to use")

    parser.add_argument('--cpus', type=int, default=None,
                        help="Number of CPUs to use")

    parser.add_argument('--number', type=int, default=maxint,
                        help="Number of test cases to analyze")

    parser.add_argument('--features', type=str, default=None,
                        help="")

    parser.add_argument('--extraction-timeout', type=int, default=300,
                        help="Timeout of the feature extraction")

    parser.add_argument('--exploration-timeout', type=int, default=600,
                        help="Timeout of the test case exploration")

    parser.add_argument('--options', type=str, default="afl,mcore,grr",
                        help="")

    parser.add_argument('--storage', type=str, default="ceo",
                        help="A folder name to store the data.")

    parser.add_argument('--blacklist', type=str, default=None,
                        help="A list of targets to ignore.")

    parser.add_argument('--verbose', type=int, default=0,
                        help="Verbosity level")

    parser.add_argument('targets', type=str,
                        help="Target file containing one executable program per line")

    parser.add_argument('mode', type=str,
                        help="Mode (init/collect/stats/stats)")

    parsed = parser.parse_args()
    return parsed

if __name__ == '__main__':


    args = parse_arguments()
    all_options = ["afl", "mcore", "grr"]
    mode = args.mode
    target_filename = args.targets
    target_blacklist = args.blacklist

    storage_dir = args.storage
    verbosity = args.verbose
    myseed = args.seed

    seed(myseed)

    # resource limits
    extraction_timeout = args.extraction_timeout
    exploration_timeout = args.exploration_timeout
    cpus = args.cpus
    nsamples = args.number

    if nsamples <= 0:
        print "[-] Invalid number of test cases to analyze."
        exit(1)

    options = []
    for option in args.options.split(","):
        if option in all_options+["none"]:
            options.append(option)
        else:
            print "[-] option", repr(option),"is invalid."
            exit(1)

    features = []
    if args.features is None or args.features == "all":
        features = features_list
    else:
        for feature in args.features.split(","):
            if feature in features_list:
                features.append(feature)
            else:
                print "[-] feature", repr(feature),"is invalid."
                exit(1)

    #print "[+] Testcases(s) loaded."
    #print "[+] Initilization storage at "+storage_dir

    if mode == "init":
        init(options, target_filename, target_blacklist, cpus, storage=storage_dir, verbose=verbosity)
        exit(0)

    if mode == "test":
        test(options, target_filename, cpus, extraction_timeout, storage=storage_dir, verbose=verbosity)
        exit(0)

    if mode == "stats":
        stats(options, features, target_filename, cpus, storage=storage_dir, verbose=verbosity)
        exit(0)

    if mode == "collect":

        names, paths = load_targets(target_filename, target_blacklist)
        print "[+] Initilization vectorizer(s) ..."
        policy = Storage(storage_dir, features_list, options)
        print "[+] Collecting data to train"

        for _ in xrange(nsamples):
            name, path = choice(zip(names, paths))
            print "[+] Picked up", name

            if "afl" in options:
                reduce_inputs(name+"/inputs", path, verbose=verbosity)

            testcases = load_testcases([name], [path])

            if len(testcases) == 0:
                print "[-] No test cases loaded for", name
                continue

            tc = choice(testcases)

            if "afl" in options:
                tc_min, label = reduce_testcase(tc, verbose=verbosity)
            else:
                tc_min = tc

            if len(tc_min) == 0:
                print "[+] Discarding empty test case"
                continue

            print "[+] Randomly selected test case", tc_min, "to explore"

            if tc_min in policy:
                print str(tc_min),"exec features already there!"
                continue
            else:
                exec_features = get_features(tc_min, extraction_timeout)

            labels, param_features = explore(options, tc_min, exploration_timeout, verbose=verbosity)
            print tc_min, labels, param_features
            policy.add(tc_min, exec_features, param_features, labels)

    else:
        print "[-] Not implemented or invalid mode!"