Selaa lähdekoodia

Added a seed-parameter for deterministic output

Denis Waßmann 7 vuotta sitten
vanhempi
commit
060ca97a31
1 muutettua tiedostoa jossa 13 lisäystä ja 1 poistoa
  1. 13 1
      code/CLI.py

+ 13 - 1
code/CLI.py

@@ -1,6 +1,7 @@
 #! /usr/bin/env python3
 import argparse
 import sys
+import random
 
 from ID2TLib.Controller import Controller
 
@@ -61,6 +62,7 @@ class CLI(object):
                             help='query the statistics database. If no query is provided, the application enters query mode.')
         parser.add_argument('-t', '--extraTests', help='perform extra tests on the input pcap file, including calculating IP entropy'
                                                        'in interval-wise, TCP checksum, and checking payload availability.', action='store_true')
+        parser.add_argument('--seed', help='random seed for the program, call the program with the same seed to get the same output')
 
 
         # Attack arguments
@@ -76,12 +78,22 @@ class CLI(object):
         """
         Decide what to do with each  of the command line parameters.
         """
+        if self.args.seed is not None:
+            self.seed_rng(self.args.seed)
+        
         if self.args.list_attacks:
             # User wants to see the available attacks
             self.process_attack_listing()
         else:
             # User wants to process a PCAP
             self.process_pcap()
+    
+    def seed_rng(self, seed):
+        try: # try to convert the seed to int
+            seed = int(seed)
+        except: pass # otherwise use the strings hash, random.seed does this automatically
+        
+        random.seed(seed)
 
     def process_attack_listing(self):
         import pkgutil
@@ -165,4 +177,4 @@ def main(args):
 
 # Uncomment to enable calling by terminal
 if __name__ == '__main__':
-    main(sys.argv[1:])
+    main(sys.argv[1:])