Multithreading Batch Processing Framework

 1 #!/usr/bin/env python
 2 # -*- coding: utf-8 -*-
 3 # Author: f0rsaken
 4
 5 import argparse
 6 import importlib
 7 import sys
 8 import threadpool
 9 import time
10
11 def main():
12     parser = argparse.ArgumentParser(description="Multithreading Batch Processing Framework")
13     parser.add_argument("-p", "--plugin", type=open, required=True, help="plugin file")
14     parser.add_argument("-i", "--input", type=open, required=True, help="input file")
15     parser.add_argument("-t", "--thread", default=1, type=int, help="number of threads")
16
17     try:
18         args = parser.parse_args()
19     except FileNotFoundError as e:
20         print(e)
21         sys.exit(1)
22
23     poolsize = args.thread
24     some_callable = get_some_callable(args.plugin)
25     list_of_args = get_list_of_args(args.input)
26
27     global log
28     log = str(int(time.time())) + ".txt"
29
30     pool = threadpool.ThreadPool(poolsize)
31     reqs = threadpool.makeRequests(some_callable, list_of_args, callback)
32     [pool.putRequest(req) for req in reqs]
33     pool.wait()
34     pool.dismissWorkers(poolsize)
35     pool.joinAllDismissedWorkers()
36
37 def get_some_callable(plugin):
38     try:
39         plugin = importlib.import_module(plugin.name.split(".")[0])
40         some_callable = plugin.exploit
41     except (AttributeError, ImportError) as e:
42         print(e)
43         sys.exit(1)
44     else:
45         return some_callable
46
47 def get_list_of_args(input):
48     list_of_args = list()
49     list_of_temp = input.readlines()
50     for i in list_of_temp:
51         list_of_args.append(i.strip())
52     return list_of_args
53
54 def callback(request, result):
55     if result:
56         with open(log, "a") as f:
57             f.write(result + "\n")
58         print(result)
59
60 if __name__ == "__main__":
61     main()