Pidfile into class and changes for this

- The PUBKEY structs from defines.py were moved to misc.py.
  This was done because having any imports in defines caused circular
  imports when using anything from defines in a class imported in this
  file.

cuckoo/apps/api.py

@@ -21,7 +21,7 @@ from cuckoo.core.database import Database, Task
 from cuckoo.core.database import TASK_REPORTED, TASK_COMPLETED, TASK_RUNNING
 from cuckoo.core.rooter import rooter
 from cuckoo.core.submit import SubmitManager
-from cuckoo.misc import cwd, version, decide_cwd, get_active_pids
+from cuckoo.misc import cwd, version, decide_cwd, Pidfile
 
 db = Database()
 sm = SubmitManager()
@@ -583,7 +583,7 @@ def cuckoo_status():
         memory=memory,
         memavail=memavail,
         memtotal=memtotal,
-        processes=get_active_pids()
+        processes=Pidfile.get_active_pids()
     )
 
     return jsonify(response)

cuckoo/common/defines.py

@@ -3,10 +3,6 @@
 # This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
 # See the file 'docs/LICENSE' for copying permission.
 
-import ctypes as c
-
-from cuckoo.misc import Structure
-
 REG_NONE = 0
 REG_SZ = 1
 REG_EXPAND_SZ = 2
@@ -15,19 +11,11 @@ REG_DWORD_LITTLE_ENDIAN = 4
 REG_DWORD = 4
 REG_DWORD_BIG_ENDIAN = 5
 
-class PUBLICKEYSTRUC(Structure):
-    _pack_ = 1
-    _fields_ = [
-        ("type", c.c_ubyte),
-        ("version", c.c_ubyte),
-        ("reserved", c.c_ushort),
-        ("algid", c.c_uint),
-    ]
+# Windows specific process rights
+# https://msdn.microsoft.com/en-us/library/ms684880(v=vs.85).aspx
+WIN_PROCESS_QUERY_INFORMATION = 0x0400
 
-class RSAPUBKEY(Structure):
-    _pack_ = 1
-    _fields_ = [
-        ("magic", c.c_uint),
-        ("bitlen", c.c_uint),
-        ("pubexp", c.c_uint),
-    ]
+# Windows specific error codes
+# https://msdn.microsoft.com/en-us/library/windows/desktop
+# /ms683189(v=vs.85).aspx
+WIN_ERR_STILL_ALIVE = 259

cuckoo/core/startup.py

@@ -17,7 +17,7 @@ import cuckoo
 from cuckoo.common.colors import red, green, yellow
 from cuckoo.common.config import Config, config, config2
 from cuckoo.common.exceptions import (
-    CuckooStartupError, CuckooFeedbackError, CuckooProcessExistsError
+    CuckooStartupError, CuckooFeedbackError
 )
 from cuckoo.common.objects import File
 from cuckoo.core.database import (
@@ -27,9 +27,7 @@ from cuckoo.core.feedback import CuckooFeedbackObject
 from cuckoo.core.log import init_logger
 from cuckoo.core.plugins import RunSignatures
 from cuckoo.core.rooter import rooter
-from cuckoo.misc import (
-    cwd, version, pidfile_exists, pid_exists, create_pidfile
-)
+from cuckoo.misc import cwd, version
 
 log = logging.getLogger(__name__)
 
@@ -408,14 +406,3 @@ def init_routing():
         if config("routing:routing:auto_rt"):
             rooter("flush_rttable", rt_table)
             rooter("init_rttable", rt_table, interface)
-
-def check_create_pidfile(name, create=True):
-    """Check if pidfile for given name exists and see if pid is still
-    alive. Create pidfile for name if it does not exist. Raises
-    CuckooProcessExistsError if pid is still alive"""
-    pid = pidfile_exists(name)
-    if pid and pid_exists(pid):
-        raise CuckooProcessExistsError("Process exists. PID: %s"
-                                       "" % pid, pid)
-    elif create:
-        create_pidfile(name)

cuckoo/main.py

@@ -32,11 +32,11 @@ from cuckoo.core.scheduler import Scheduler
 from cuckoo.core.startup import (
     check_configs, init_modules, check_version, init_logfile, init_logging,
     init_console_logging, init_tasks, init_yara, init_binaries, init_rooter,
-    init_routing, check_create_pidfile
+    init_routing
 )
 from cuckoo.misc import (
     cwd, load_signatures, getuser, decide_cwd, drop_privileges, is_windows,
-    remove_pidfile
+    Pidfile, pid_exists
 )
 
 log = logging.getLogger("cuckoo")
@@ -108,7 +108,11 @@ def cuckoo_init(level, ctx, cfg=None):
 
     # Only one Cuckoo process should exist per CWD.
     # Run this check before any files are possibly modified
-    check_create_pidfile("cuckoo")
+    pidf = Pidfile("cuckoo")
+    if pidf.exists() and pid_exists(pidf.pid):
+        raise CuckooProcessExistsError("", pidf.pid)
+    else:
+        pidf.create()
 
     init_console_logging(level)
 
@@ -175,7 +179,7 @@ def cuckoo_main(max_analysis_count=0):
         sched.start()
     except KeyboardInterrupt:
         sched.stop()
-        remove_pidfile("cuckoo")
+    Pidfile("cuckoo").remove()
 
 @click.group(invoke_without_command=True)
 @click.option("-d", "--debug", is_flag=True, help="Enable verbose logging")
@@ -338,13 +342,14 @@ def process(ctx, instance, report, maxcount):
     init_console_logging(level=ctx.parent.level)
 
     if instance:
-        try:
-            check_create_pidfile(instance)
-        except CuckooProcessExistsError as e:
+        pidf = Pidfile(instance)
+        if pidf.exists() and pid_exists(pidf.pid):
             message = "Cuckoo process instance \'%s\' already exists. " \
-                      "PID: %s\n" % (instance, e.pid)
-            log.error(message)
+                      "PID: %s\n" % (instance, pidf.pid)
+            log.error(red(message))
             sys.exit(1)
+        else:
+            pidf.create()
 
         init_logfile("process-%s.json" % instance)
 
@@ -380,8 +385,9 @@ def process(ctx, instance, report, maxcount):
             process_tasks(instance, maxcount)
     except KeyboardInterrupt:
         print(red("Aborting (re-)processing of your analyses.."))
-        if instance:
-            remove_pidfile(instance)
+
+    if instance:
+        Pidfile(instance).remove()
 
 @main.command()
 @click.argument("socket", type=click.Path(readable=False, dir_okay=False), default="/tmp/cuckoo-rooter", required=False)

cuckoo/misc.py

@@ -21,8 +21,12 @@ except ImportError:
 
 import cuckoo
 
+from cuckoo.common.defines import (
+    WIN_PROCESS_QUERY_INFORMATION, WIN_ERR_STILL_ALIVE
+)
 from cuckoo.common.exceptions import CuckooStartupError
 
+
 log = logging.getLogger(__name__)
 
 # Cuckoo Working Directory base path.
@@ -237,16 +241,11 @@ def pid_exists(pid):
     if is_windows():
         from ctypes import windll, wintypes
         dw_exit = wintypes.DWORD()
-        # https://msdn.microsoft.com/en-us/library/ms684880(v=vs.85).aspx
-        # 1024 = PROCESS_QUERY_INFORMATION
-        proc_h = windll.kernel32.OpenProcess(1024, 0, pid)
+        proc_h = windll.kernel32.OpenProcess(WIN_PROCESS_QUERY_INFORMATION,
+                                             0, pid)
         windll.kernel32.GetExitCodeProcess(proc_h, ctypes.byref(dw_exit))
 
-        # https://msdn.microsoft.com/en-us/library/windows/desktop
-        # /ms683189(v=vs.85).aspx
-        # 259 = STILL_ACTIVE
-
-        return dw_exit.value == 259
+        return dw_exit.value == WIN_ERR_STILL_ALIVE
 
     elif is_linux() or is_macosx():
         # Send signal 0 to process. Exception will be thrown if it does not
@@ -262,60 +261,92 @@ def pid_exists(pid):
         # Return none if not of of the above platforms
         return None
 
-def create_pidfile(name):
-    """Creates a file with a .pid extension containing
-    the pid of the process executing this method
-    @param name: name without extension to use for file
-    """
-    pid_path = cwd("pidfiles")
-    if not os.path.exists(pid_path):
-        os.mkdir(pid_path)
+class Pidfile(object):
 
-    with open(os.path.join(pid_path, "%s.pid"% name), "wb") as fw:
-        fw.write(str(os.getpid()))
+    def __init__(self, name=None):
+        """Manage pidfile of given name
+        @param name: name without extension to use for file
+        """
+        self.dir = cwd("pidfiles")
+        self.name = None
+        self.pidfile_path = None
+        self.pid = None
+        if name:
+            self.set_name(name)
 
-def remove_pidfile(name):
-    """Remove pidfile of name if it exists
-    @param name: name without extension to use for file
-    """
-    pid_path = cwd("pidfiles", "%s.pid" % name)
-    if os.path.exists(pid_path):
-        os.remove(pid_path)
+    def set_name(self, name):
+        """Change pidfile name"""
+        if not name.endswith(".pid"):
+            name = "%s.pid" % name
+        self.name = name
+        self.pidfile_path = os.path.join(self.dir, self.name)
 
-def pidfile_exists(name):
-    """Check if a pid file exists for the given name. Returns int PID
-    if the pidfile exists, False if not. None if invalid pidfile
-    @param name: name without extension to use for file
-    """
-    if not name.endswith(".pid"):
-        name = "%s.pid" % name
+    def create(self):
+        """Creates a file with a .pid extension containing
+        the pid of the process executing this method"""
+        if not os.path.exists(self.dir):
+            os.mkdir(self.dir)
 
-    pidfile_path = cwd("pidfiles", name)
-    if not os.path.exists(pidfile_path):
-        return False
+        with open(self.pidfile_path, "wb") as fw:
+            fw.write(str(os.getpid()))
 
-    with open(pidfile_path, "rb") as fp:
-        try:
-            pid = int(fp.read())
-        except ValueError:
-            # Return None if pidfile contains invalid value
-            return None
+    def remove(self):
+        """Remove pidfile of name if it exists"""
+        if os.path.exists(self.pidfile_path):
+            os.remove(self.pidfile_path)
+            return True
+        else:
+            return False
 
-    return pid
+    def exists(self):
+        """Check if a pid file exists for the given name. Returns int PID
+        if the pidfile exists, False if not. None if invalid pidfile"""
+        if not os.path.exists(self.pidfile_path):
+            return False
 
-def get_active_pids():
-    """Return a dict containing active pids.
-    Key is the pidfile name and value is pid"""
+        return self.read()
+
+    def read(self):
+        """Read PID from pidfile. Will raise IOError if pidfile does
+        not exist"""
+        with open(self.pidfile_path, "rb") as fp:
+            try:
+                self.pid = int(fp.read())
+            except ValueError:
+                # Set pid to None if pidfile contains invalid value
+                self.pid = None
+
+        return self.pid
+
+    @staticmethod
+    def get_active_pids():
+        """Return a dict containing active pids.
+        Key is the pidfile name and value is pid"""
+        pids = {}
+
+        pidfile = Pidfile()
+        for pidf in os.listdir(pidfile.dir):
+            pidfile.set_name(pidf)
+            pid = pidfile.read()
+            if pid:
+                if pid_exists(pid):
+                    pids[pidf.split(".", 1)[0]] = pid
 
-    pidfile_path = cwd("pidfiles")
-    pids = {}
-    if not os.path.exists(pidfile_path):
         return pids
 
-    for pidf in os.listdir(pidfile_path):
-        pid = pidfile_exists(pidf)
-        if pid:
-            if pid_exists(pid):
-                pids[pidf.split(".",1)[0]] = pid
+class PUBLICKEYSTRUC(Structure):
+    _pack_ = 1
+    _fields_ = [
+        ("type", ctypes.c_ubyte),
+        ("version", ctypes.c_ubyte),
+        ("reserved", ctypes.c_ushort),
+        ("algid", ctypes.c_uint),
+    ]
 
-    return pids
+class RSAPUBKEY(Structure):
+    _pack_ = 1
+    _fields_ = [
+        ("magic", ctypes.c_uint),
+        ("bitlen", ctypes.c_uint),
+        ("pubexp", ctypes.c_uint),
+    ]

tests/test_defines.py

@@ -2,18 +2,19 @@
 # This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
 # See the file 'docs/LICENSE' for copying permission.
 
-from cuckoo.common.defines import PUBLICKEYSTRUC, RSAPUBKEY, REG_DWORD
+from cuckoo.common.defines import (
+    REG_NONE, REG_SZ, REG_EXPAND_SZ, REG_BINARY, REG_DWORD_LITTLE_ENDIAN,
+    REG_DWORD, REG_DWORD_BIG_ENDIAN, WIN_PROCESS_QUERY_INFORMATION,
+    WIN_ERR_STILL_ALIVE
+)
 
 def test_defines():
-    a = PUBLICKEYSTRUC.from_buffer_copy("A"*8)
-    assert a.type == 0x41
-    assert a.version == 0x41
-    assert a.reserved == 0x4141
-    assert a.algid == 0x41414141
-
-    a = RSAPUBKEY.from_buffer_copy("A"*12)
-    assert a.magic == 0x41414141
-    assert a.bitlen == 0x41414141
-    assert a.pubexp == 0x41414141
-
+    assert REG_NONE == 0
+    assert REG_SZ == 1
+    assert REG_EXPAND_SZ == 2
+    assert REG_BINARY == 3
+    assert REG_DWORD_LITTLE_ENDIAN == 4
     assert REG_DWORD == 4
+    assert REG_DWORD_BIG_ENDIAN == 5
+    assert WIN_PROCESS_QUERY_INFORMATION == 0x0400
+    assert WIN_ERR_STILL_ALIVE == 259

tests/test_startup.py

@@ -19,10 +19,10 @@ from cuckoo.common.objects import File
 from cuckoo.core.database import Database
 from cuckoo.core.startup import (
     init_modules, check_version, init_rooter, init_routing, init_yara,
-    init_tasks, init_binaries, check_create_pidfile
+    init_tasks, init_binaries
 )
 from cuckoo.main import cuckoo_create
-from cuckoo.misc import set_cwd, load_signatures, cwd, is_linux, create_pidfile
+from cuckoo.misc import set_cwd, load_signatures, cwd, is_linux
 
 def test_init_tasks():
     def init(reschedule):
@@ -644,52 +644,3 @@ class TestYaraIntegration(object):
         init_yara()
         assert len(File.yara_rules) == 5
         assert not list(File.yara_rules["binaries"])
-
-def test_create_check_pidfile_exists():
-    set_cwd(tempfile.mkdtemp())
-    create_pidfile("test")
-
-    with pytest.raises(CuckooProcessExistsError) as e:
-        check_create_pidfile("test")
-        assert e.pid == os.getpid()
-
-def test_create_check_pidfile_dead_proc():
-    set_cwd(tempfile.mkdtemp())
-    pidfiles_path = cwd("pidfiles")
-    pidfile = cwd("pidfiles", "test2.pid")
-    os.mkdir(pidfiles_path)
-
-    with open(pidfile, "wb") as fw:
-        fw.write("9991337")
-
-    try:
-        check_create_pidfile("test2")
-    except CuckooProcessExistsError:
-        pytest.fail("Raised exception. It should not do so if the pid"
-                    " does not exist.")
-
-    with open(pidfile, "rb") as fp:
-        assert int(fp.read()) == os.getpid()
-
-def test_create_check_pidfile():
-    set_cwd(tempfile.mkdtemp())
-    try:
-        check_create_pidfile("test3")
-    except CuckooProcessExistsError:
-        pytest.fail("Raised exception. It should not do so if the pidfile "
-                    "and process do not exist")
-
-    pidfile = cwd("pidfiles", "test3.pid")
-    assert os.path.exists(pidfile)
-    with open(pidfile, "rb") as fp:
-        assert int(fp.read()) == os.getpid()
-
-def test_create_check_pidfile_createfalse():
-    set_cwd(tempfile.mkdtemp())
-    try:
-        check_create_pidfile("test4", create=False)
-    except CuckooProcessExistsError:
-        pytest.fail("Raised exception. It should not do so if the pidfile "
-                    "and process do not exist")
-
-    assert not os.path.exists(cwd("pidfiles", "test4.pid"))