improved pidfile handling

c0m4r · Jan 1, 2024 · 882e922 · 882e922
1 parent 06e5520
commit 882e922
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 25 deletions.
diff --git a/loki-daemonized.patch b/loki-daemonized.patch
@@ -1,5 +1,5 @@
 --- loki.py.original	2023-12-28 14:31:22.375121479 +0100
-+++ loki.py	2024-01-01 23:30:04.789802470 +0100
++++ loki.py	2024-01-01 23:57:18.557125199 +0100
 @@ -48,6 +48,10 @@
  from lib.doublepulsar import DoublePulsar
  from lib.vuln_checker import VulnChecker
@@ -144,22 +144,33 @@
 
          except Exception:
              if logger.debug:
-@@ -1452,6 +1496,14 @@
-         print('LOKI\'s work has been interrupted by a human. Returning to Asgard.')
-     sys.exit(0)
+@@ -1444,12 +1488,25 @@
 
-+def signal_handler_term(signal_name, frame):
-+    try:
-+        os.remove(args.pidfile)
-+    except:
-+        pass
-+    print('SIGTERM')
+
+ # CTRL+C Handler --------------------------------------------------------------
++def remove_pidfile():
++    if(args.d):
++        try:
++            os.remove(args.pidfile)
++        except Exception:
++            pass
++
+ def signal_handler(signal_name, frame):
+     try:
+         print("------------------------------------------------------------------------------\n")
+         logger.log('INFO', 'Init', 'LOKI\'s work has been interrupted by a human. Returning to Asgard.')
+     except Exception:
+         print('LOKI\'s work has been interrupted by a human. Returning to Asgard.')
++    remove_pidfile()
 +    sys.exit(0)
 +
++def signal_handler_term(signal_name, frame):
++    remove_pidfile()
++    print('LOKI\'s work has been interrupted by a SIGTERM. Returning to Asgard.')
+     sys.exit(0)
+
  def main():
-     """
-     Argument parsing function
-@@ -1468,6 +1520,12 @@
+@@ -1468,6 +1525,12 @@
      parser.add_argument('-a', help='Alert score', metavar='alert-level', default=100)
      parser.add_argument('-w', help='Warning score', metavar='warning-level', default=60)
      parser.add_argument('-n', help='Notice score', metavar='notice-level', default=40)
@@ -172,7 +183,7 @@
      parser.add_argument('--allhds', action='store_true', help='Scan all local hard drives (Windows only)', default=False)
      parser.add_argument('--alldrives', action='store_true', help='Scan all drives (including network drives and removable media)', default=False)
      parser.add_argument('--printall', action='store_true', help='Print all files that are scanned', default=False)
-@@ -1532,9 +1590,17 @@
+@@ -1532,9 +1595,25 @@
      # Signal handler for CTRL+C
      signal_module.signal(signal_module.SIGINT, signal_handler)
 
@@ -184,13 +195,21 @@
 
 +    # Save pidfile
 +    if args.d is True:
-+        with open(args.pidfile, 'w', encoding='utf-8') as f:
-+            f.write(str(os.getpid()))
++        if(os.path.exists(args.pidfile)):
++            fpid = open(args.pidfile, 'r')
++            loki_pid = int(fpid.read())
++            fpid.close()
++            if psutil.pid_exists(loki_pid):
++                print("LOKI daemon already running. Returning to Asgard.")
++                sys.exit(0)
++        with open(args.pidfile, 'w', encoding='utf-8') as fpid:
++            fpid.write(str(os.getpid()))
++            fpid.close()
 +
      # Remove old log file
      if os.path.exists(args.l):
          os.remove(args.l)
-@@ -1553,8 +1619,17 @@
+@@ -1553,8 +1632,17 @@
          updateLoki(sigsOnly=False)
          sys.exit(0)
 
@@ -209,7 +228,7 @@
 
      # Loki
      loki = Loki(args.intense)
-@@ -1619,6 +1694,63 @@
+@@ -1619,6 +1707,63 @@
                  loki.scan_path(defaultPath)
 
          # Linux & macOS

diff --git a/loki.py b/loki.py
@@ -1488,20 +1488,25 @@ def walk_error(err):
 
 
 # CTRL+C Handler --------------------------------------------------------------
+def remove_pidfile():
+    if(args.d):
+        try:
+            os.remove(args.pidfile)
+        except Exception:
+            pass
+
 def signal_handler(signal_name, frame):
     try:
         print("------------------------------------------------------------------------------\n")
         logger.log('INFO', 'Init', 'LOKI\'s work has been interrupted by a human. Returning to Asgard.')
     except Exception:
         print('LOKI\'s work has been interrupted by a human. Returning to Asgard.')
+    remove_pidfile()
     sys.exit(0)
 
 def signal_handler_term(signal_name, frame):
-    try:
-        os.remove(args.pidfile)
-    except:
-        pass
-    print('SIGTERM')
+    remove_pidfile()
+    print('LOKI\'s work has been interrupted by a SIGTERM. Returning to Asgard.')
     sys.exit(0)
 
 def main():
@@ -1598,8 +1603,16 @@ def main():
 
     # Save pidfile
     if args.d is True:
-        with open(args.pidfile, 'w', encoding='utf-8') as f:
-            f.write(str(os.getpid()))
+        if(os.path.exists(args.pidfile)):
+            fpid = open(args.pidfile, 'r')
+            loki_pid = int(fpid.read())
+            fpid.close()
+            if psutil.pid_exists(loki_pid):
+                print("LOKI daemon already running. Returning to Asgard.")
+                sys.exit(0)
+        with open(args.pidfile, 'w', encoding='utf-8') as fpid:
+            fpid.write(str(os.getpid()))
+            fpid.close()
 
     # Remove old log file
     if os.path.exists(args.l):