From 7c15dcf0b426c4d91f6315924dec133b720596bb Mon Sep 17 00:00:00 2001
From: karimpwnz <karim@karimrahal.com>
Date: Fri, 16 Aug 2019 16:52:24 +0300
Subject: [PATCH] Introduce _random_ feature as per #50

---
 Interlace/lib/core/input.py | 37 +++++++++++++++++++++++++++++++++----
 1 file changed, 33 insertions(+), 4 deletions(-)

diff --git a/Interlace/lib/core/input.py b/Interlace/lib/core/input.py
index 4921ddd..ca3cfec 100644
--- a/Interlace/lib/core/input.py
+++ b/Interlace/lib/core/input.py
@@ -7,16 +7,22 @@ import sys
 from re import compile
 from random import sample
 from math import ceil
+import random
 
 
 class InputHelper(object):
     @staticmethod
-    def readable_file(parser, arg):
+    def check_path(parser, arg):
         if not os.path.exists(arg):
-            parser.error("The file %s does not exist!" % arg)
+            parser.error("The path %s does not exist!" % arg)
         else:
-            return open(arg, 'r')  # return an open file handle
+            return arg
 
+    @staticmethod
+    def readable_file(parser, arg):
+        if InputHelper.check_path(parser, arg):
+            return open(arg, 'r')  # return an open file handle
+            
     @staticmethod
     def check_positive(parser, arg):
         ivalue = int(arg)
@@ -25,6 +31,17 @@ class InputHelper(object):
 
         return arg
 
+    @staticmethod
+    def _get_files_from_directory(arg):
+        files = list()
+
+        for file in os.listdir(arg):
+            location = os.path.join(arg, file)
+            if os.path.isfile(location):
+                files.append(location)
+
+        return files
+
     @staticmethod
     def _get_ips_from_range(ip_range):
         ips = set()
@@ -128,7 +145,6 @@ class InputHelper(object):
             else:
                 real_ports = [arguments.realport]
 
-
         # process targets first
         if arguments.target:
             ranges.add(arguments.target)
@@ -197,6 +213,10 @@ class InputHelper(object):
         if len(targets) == 0:
             raise Exception("No target provided, or empty target list")
 
+        if arguments.random:
+            files = InputHelper._get_files_from_directory(arguments.random)
+            random_file = random.choice(files)
+
         if arguments.command:
             commands.add(arguments.command.rstrip('\n'))
         else:
@@ -211,6 +231,9 @@ class InputHelper(object):
 
         if arguments.realport:
             final_commands = InputHelper._replace_variable_for_commands(final_commands, "_realport_", real_ports)
+        
+        if arguments.random:
+            final_commands = InputHelper._replace_variable_for_commands(final_commands, "_random_", [random_file])
 
         if arguments.output:
             final_commands = InputHelper._replace_variable_for_commands(final_commands, "_output_", [arguments.output])
@@ -338,6 +361,12 @@ class InputParser(object):
             help='Specify a real port variable that can be used in commands as _realport_'
         )
 
+        parser.add_argument(
+            '-random', dest='random',
+            help='Specify a directory of files that can be randomly used in commands as _random_',
+            type=lambda x: InputHelper.check_path(parser, x)
+        )
+
         parser.add_argument(
             '--no-cidr', dest='nocidr', action='store_true', default=False,
             help='If set then CIDR notation in a target file will not be automatically '