LOKI upgrader allows a signature clean-up (--clean)

Neo23x0 · Apr 14, 2018 · 6c68ad4 · 6c68ad4
1 parent fc2f405
commit 6c68ad4
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/loki-upgrader.py b/loki-upgrader.py
@@ -55,7 +55,7 @@ def __init__(self, debug, logger, application_path):
         self.logger = logger
         self.application_path = application_path
 
-    def update_signatures(self):
+    def update_signatures(self, clean=False):
         try:
             for sig_url in self.UPDATE_URL_SIGS:
                 # Downloading current repository
@@ -72,6 +72,9 @@ def update_signatures(self):
                 # Preparations
                 try:
                     sigDir = os.path.join(self.application_path, './signature-base/')
+                    if clean:
+                        self.logger.log("INFO", "Upgrader", "Cleaning directory '%s'" % sigDir)
+                        shutil.rmtree(sigDir)
                     for outDir in ['', 'iocs', 'yara', 'misc']:
                         fullOutDir = os.path.join(sigDir, outDir)
                         if not os.path.exists(fullOutDir):
@@ -208,6 +211,8 @@ def get_application_path():
     parser.add_argument('--progonly', action='store_true', help='Update the program files only', default=False)
     parser.add_argument('--nolog', action='store_true', help='Don\'t write a local log file', default=False)
     parser.add_argument('--debug', action='store_true', default=False, help='Debug output')
+    parser.add_argument('--clean', action='store_true', default=False, help='Clean up the signature directory and get '
+                                                                            'a fresh set')
     parser.add_argument('--detached', action='store_true', default=False, help=argparse.SUPPRESS)
 
     args = parser.parse_args()
@@ -221,16 +226,15 @@ def get_application_path():
     # Logger
     logger = LokiLogger(args.nolog, args.l, t_hostname, '', '', False, False, args.debug, platform=platform, caller='upgrader')
 
-    # Update Loki
+    # Update LOKI
     updater = LOKIUpdater(args.debug, logger, get_application_path())
 
-    # Updating LOKI
     if not args.sigsonly:
         logger.log("INFO", "Upgrader", "Updating LOKI ...")
         updater.update_loki()
     if not args.progonly:
         logger.log("INFO", "Upgrader", "Updating Signatures ...")
-        updater.update_signatures()
+        updater.update_signatures(args.clean)
 
     logger.log("INFO", "Upgrader", "Update complete")