TSScraper/main.py

113 lines
3.3 KiB
Python

"""
main.py
"""
import re
import os
import os.path
import sys
import multiprocessing
import importlib
import os.path
import timeit
import cProfile
import tsscraper
class Application(object):
thread_count = 8
threads = None
target_directory = None
target_exporter = None
def print_usage(self):
print("Usage: '%s <exporter> <output directory> <target directories...>'" % sys.argv[0])
print("Or: '%s exporters' for a list of known exporters." % sys.argv[0])
def get_available_exporters(self):
exporters = { }
for root, dirs, files in os.walk("exporters"):
for filename in files:
module_name, extension = os.path.splitext(filename)
if (module_name == "__init__"):
continue
try:
module = importlib.import_module('exporters.%s' % (module_name))
exporters[module_name] = module
except ImportError as e:
print(e)
return exporters
def main(self):
"""
The main entry point of the application. This is equivalent to
the main() method in C and C++.
"""
if (len(sys.argv) < 2):
self.print_usage()
return
exporters = self.get_available_exporters()
if (sys.argv[1] == "exporters"):
print("Available Exporters: ")
for exporter in exporters:
print("\t- %s" % exporter)
print("\t- None")
return
elif(len(sys.argv) < 4):
self.print_usage()
return
self.target_directory = sys.argv[3]
self.output_directory = sys.argv[2]
self.target_exporter = sys.argv[1]
self.run()
def run(self):
exporter = None
if (self.target_exporter.lower() != "none"):
exporters = self.get_available_exporters()
try:
exporter = exporters[self.target_exporter]
except KeyError as e:
print("Error: No such exporter '%s'." % self.target_exporter)
self.print_usage()
return
# First, process base
base_results = None
if (os.path.isdir("base") is False):
print("Warning: No local copy of base found! Some reference checks will report false positives.")
else:
print("INFO: Processing base ...")
base_scraper = tsscraper.TSScraper("base", self.thread_count)
base_results = base_scraper.process()
print("INFO: Processing '%s' ..." % self.target_directory)
scraper = tsscraper.TSScraper(self.target_directory, self.thread_count, base_results)
results = scraper.process()
# Init the exporter
print("INFO: Exporting data ...")
if (exporter is not None):
# Ensure that the output directory at least exists
try:
os.mkdir(self.output_directory)
except OSError:
pass
output = exporter.Exporter(results, self.target_directory)
output.write(self.output_directory)
if __name__ == "__main__":
print("Operation Completion-----------------------\n%f Seconds" % timeit.timeit("Application().main()", number=1, setup="from __main__ import Application"))