|
|
@@ -17,6 +17,13 @@ def parse_args():
|
|
|
parser.add_argument('-s', '--sample-size', type=int,
|
|
|
default="200", help='number of packets in a sample \
|
|
|
(default: 200)')
|
|
|
+ parser.add_argument('-r', '--randomize', action='store_const', default=False,
|
|
|
+ const=True, help='Pre-randomize sample order')
|
|
|
+ parser.add_argument('-f', '--fix-sample-count', action='store_const',
|
|
|
+ default=False, const=True,
|
|
|
+ help='Keep all users as a fixed sample count.')
|
|
|
+ parser.add_argument('-m', '--min', type=int, default="5",
|
|
|
+ help='Minimum number of samples per user (default: 5)')
|
|
|
parser.add_argument('-l', '--low-act-threshold', type=float, default=1.0,
|
|
|
help='P/s below which is considered low activity \
|
|
|
(default: 1.0 P/s)')
|
|
|
@@ -41,13 +48,33 @@ def main():
|
|
|
Sample.set_activity_thresholds(args.low_act_threshold,
|
|
|
args.high_act_threshold,
|
|
|
args.lookback)
|
|
|
+ users = {}
|
|
|
+ for line in args.match_file:
|
|
|
+ if "pcap" in line:
|
|
|
+ samples = list(Sample.make_samples(*line.split(" "), args.sample_size))
|
|
|
+ user = samples[0].user
|
|
|
+ if samples:
|
|
|
+ if user in users:
|
|
|
+ users[user].extend(samples)
|
|
|
+ else:
|
|
|
+ users[user] = samples
|
|
|
+ fix_point = 99999999999
|
|
|
+ if args.fix_sample_count:
|
|
|
+ for u in users:
|
|
|
+ print(len(users[u]))
|
|
|
+ if len(users[u]) >= args.min and len(users[u]) < fix_point:
|
|
|
+ fix_point = len(users[u])
|
|
|
+ print(fix_point)
|
|
|
out = [sample
|
|
|
- for line in args.match_file if "pcap" in line
|
|
|
- for sample in Sample.make_samples(*line.split(" "), args.sample_size)]
|
|
|
+ for u in users if len(users[u]) >= args.min
|
|
|
+ for sample in users[u][0:fix_point]]
|
|
|
try:
|
|
|
import cPickle as pickle
|
|
|
except:
|
|
|
import pickle
|
|
|
+ if args.randomize:
|
|
|
+ from random import shuffle
|
|
|
+ shuffle(out)
|
|
|
pickle.dump(out, args.outfile)
|
|
|
|
|
|
if __name__ == '__main__':
|