Merge branch 'master' of https://git.tflucke.name/tflucke/SSH-Master-Thesis

src/feature-extractor/extractor.py

@@ -1,6 +1,5 @@
 #!/usr/bin/python3
-import typing
-import sys
+import typing, sys, os
 
 def parse_args():
     import argparse
@@ -43,10 +42,12 @@ def parse_args():
     parser.add_argument('-d', '--dataset', choices=["guided", "free", "both"],
                         default="uniform", help='Which dataset to pull from. \
                         (One of guided, free, or both; default: \"both\")')
+    parser.add_argument('-u', '--min-users', type=int, default="20",
+                        help='Minimum number of users to consider a valid \
+                        dataset. (default: 20)')
     return parser.parse_args()
 
 def enter_data_dir(match_file: typing.TextIO):
-    import os
     if match_file is not sys.stdin:
         os.chdir(os.path.dirname(match_file.name))
 
@@ -71,6 +72,10 @@ def main():
                     users[user].extend(samples)
                 else:
                     users[user] = samples
+    if len(users) < args.min_users:
+        print("Not enough data to build valid dataset.", file=sys.stderr)
+        os.remove(args.outfile)
+        return
     fix_point = 99999999999
     if args.fix_sample_count:
         for u in users: