<commit_before>#!/usr/bin/env python
#
# $Id$
#

"""psutil is a module providing convenience functions for managing
processes in a portable way by using Python.
"""

__all__ = [
    "NoSuchProcess",
    "AccessDenied",
    "NUM_CPUS",
    "ProcessInfo",
    "Process",
    "test",
    "CPUTimes",
    "pid_exists",
    "get_pid_list",
    "process_iter",
    "get_process_list",
    "TOTAL_PHYMEM",
    "avail_phymem",
    "used_phymem",
    "total_virtmem",
    "avail_virtmem",
    "used_virtmem",
    "cpu_times",
    "cpu_percent",
    ]

__version__ = '0.1.3'

import sys
import os
import time
try:
    import pwd, grp
except ImportError:
    pwd = grp = None

# exceptions are imported here, but may be overriden by platform
# module implementation later
from error import *

# import the appropriate module for our platform only
if sys.platform.lower().startswith("linux"):
    from _pslinux import *

elif sys.platform.lower().startswith("win32"):
    from _psmswindows import *

elif sys.platform.lower().startswith("darwin"):
    from _psosx import *

elif sys.platform.lower().startswith("freebsd"):
    from _psbsd import *
else:
    raise ImportError('no os specific module found')

# platform-specific modules define an Impl implementation class
_platform_impl = Impl()

class CPUTimes:
    """This class contains information about CPU times.
    It is not used directly but it's returned as an instance by
    psutil.cpu_times() function.

    Every CPU time is accessible in form of an attribute and represents
    the time CPU has spent in the given mode.

    The attributes availability varies depending on the platform.
    Here follows a list of all available attributes:

     - user
     - system
     - idle
     - nice (UNIX)
     - iowait (Linux)
     - irq (Linux, FreeBSD)
     - softirq (Linux)
    """

    def __init__(self, **kwargs):
        self.__attrs = []
        for name in kwargs:
            setattr(self, name, kwargs[name])
            self.__attrs.append(name)

    def __str__(self):
        string = []
        for attr in self.__attrs:
            value = getattr(self, attr)
            string.append("%s=%s" %(attr, value))
        return '; '.join(string)

    def __iter__(self):
        for attr in self.__attrs:
            yield getattr(self, attr)

class ProcessInfo(object):
    """Class that allows the process information to be passed between
    external code and psutil.  Used directly by the Process class.
    """

    def __init__(self, pid, ppid=None, name=None, path=None, cmdline=None,
                       uid=None, gid=None):
        self.pid = pid
        self.ppid = ppid
        self.name = name
        self.path = path
        self.cmdline = cmdline
        # if we have the cmdline but not the path, figure it out from argv[0]
        if cmdline and not path:
            self.path = os.path.dirname(cmdline[0])
        self.uid = uid
        self.gid = gid
        self.create = None
        self.username = None
        self.groupname = None

class Process(object):
    """Represents an OS process."""

    def __init__(self, pid):
        """Create a new Process object, raises NoSuchProcess if the PID
        does not exist, and ValueError if the parameter is not an
        integer PID."""
        if not isinstance(pid, int):
            raise ValueError("An integer is required")
        if not pid_exists(pid):
            raise NoSuchProcess("No process found with PID %s" % pid)
        self._procinfo = ProcessInfo(pid)
        self.is_proxy = True
        # try to init CPU times, if it raises AccessDenied then suppress
        # it so it won't interrupt the constructor. First call to
        # get_cpu_percent() will trigger the AccessDenied exception
        # instead.
        try:
            self._last_sys_time = time.time()
            self._last_user_time, self._last_kern_time = self.get_cpu_times()
        except AccessDenied:
            self._last_user_time, self._last_kern_time = None, None

    def __str__(self):
        return "psutil.Process <PID:%s; PPID:%s; NAME:'%s'; PATH:'%s'; " \
               "CMDLINE:%s; UID:%s; GID:%s;>" \
               %(self.pid, self.ppid, self.name, self.path, self.cmdline, \
                 self.uid, self.gid)

    def __eq__(self, other):
        """Test for equality with another Process object based on PID,
        name etc."""
        if self.pid != other.pid:
            return False

        # make sure both objects are initialized
        self.deproxy()
        other.deproxy()
        # check all non-callable and non-private attributes for equality
        # if the attribute is missing from other then return False
        for attr in dir(self):
            attrobj = getattr(self, attr)
            # skip private attributes
            if attr.startswith("_"):
                continue

            # skip methods and Process objects
            if callable(attrobj) or isinstance(attrobj, Process):
                continue

            # attribute doesn't exist or isn't equal, so return False
            if not hasattr(other, attr):
                return False
            if getattr(self, attr) != getattr(other, attr):
                return False

        return True

    def deproxy(self):
        """Used internally by Process properties. The first call to
        deproxy() initializes the ProcessInfo object in self._procinfo
        with process data read from platform-specific module's
        get_process_info() method.

        This method becomes a NO-OP after the first property is accessed.
        Property data is filled in from the ProcessInfo object created,
        and further calls to deproxy() simply return immediately without
        calling get_process_info().
        """
        if self.is_proxy:
            # get_process_info returns a tuple we use as the arguments
            # to the ProcessInfo constructor
            self._procinfo = ProcessInfo(*_platform_impl.get_process_info(self._procinfo.pid))
            self.is_proxy = False

    @property
    def pid(self):
        """The process pid."""
        return self._procinfo.pid

    @property
    def ppid(self):
        """The process parent pid."""
        self.deproxy()
        return self._procinfo.ppid

    @property
    def parent(self):
        """Return the parent process as a Process object. If no ppid is
        known then return None."""
        if self.ppid is not None:
            return Process(self.ppid)
        return None

    @property
    def name(self):
        """The process name."""
        self.deproxy()
        return self._procinfo.name

    @property
    def path(self):
        """The process path."""
        self.deproxy()
        return self._procinfo.path

    @property
    def cmdline(self):
        """The command line process has been called with."""
        self.deproxy()
        return self._procinfo.cmdline

    @property
    def uid(self):
        """The real user id of the current process."""
        self.deproxy()
        return self._procinfo.uid

    @property
    def gid(self):
        """The real group id of the current process."""
        self.deproxy()
        return self._procinfo.gid

    @property
    def username(self):
        """The name of the user that owns the process."""
    	if self._procinfo.username is not None:
    		return self._procinfo.username
        if pwd is not None:
            self._procinfo.username = pwd.getpwuid(self.uid).pw_name
        else:
            self._procinfo.username =  _platform_impl.get_process_username(self.pid)
    	return self._procinfo.username

    @property
    def groupname(self):
        """The real groupname of the current process."""
    	if self._procinfo.groupname is not None:
    		return self._procinfo.groupname
        if grp is not None:
            self._procinfo.groupname = grp.getgrgid(self.gid).gr_name
        else:
            # XXX - provisional for Windows
            return "groupname"
    	return self._procinfo.groupname

    @property
    def environ(self):
        """The process environment variables as a dictionary."""
        return _platform_impl.get_process_environ(self.pid)

    @property
    def create_time(self):
        """The process creation time as a floating point number
        expressed in seconds since the epoch, in UTC.
        """
        if self._procinfo.create is None:
            self._procinfo.create = _platform_impl.get_process_create_time(self.pid)
        return self._procinfo.create

    def getcwd(self):
        """Return a string representing the process current working directory."""
        return _platform_impl.get_process_cwd(self.pid)

    def get_cpu_percent(self):
        """Compare process times to system time elapsed since last call
        and calculate CPU utilization as a percentage. It is recommended
        for accuracy that this function be called with at least 1 second
        between calls. The initial delta is calculated from the
        instantiation of the Process object.
        """
        now = time.time()
        # will raise AccessDenied on OS X if not root or in procmod group
        user_t, kern_t = _platform_impl.get_cpu_times(self.pid)

        total_proc_time = float((user_t - self._last_user_time) + \
                                (kern_t - self._last_kern_time))
        try:
            percent = total_proc_time / float((now - self._last_sys_time))
        except ZeroDivisionError:
            percent = 0.000

        # reset the values
        self._last_sys_time = time.time()
        self._last_user_time, self._last_kern_time = self.get_cpu_times()

        return (percent * 100.0) / NUM_CPUS

    def get_cpu_times(self):
        """Return a tuple whose values are process CPU user and system
        time.  These are the same first two values that os.times()
        returns for the current process.
        """
        return _platform_impl.get_cpu_times(self.pid)

    def get_memory_info(self):
        """Return a tuple representing RSS (Resident Set Size) and VMS
        (Virtual Memory Size) in bytes.

        On UNIX RSS and VMS are the same values shown by ps.

        On Windows RSS and VMS refer to "Mem Usage" and "VM Size" columns
        of taskmgr.exe.
        """
        return _platform_impl.get_memory_info(self.pid)

    def get_memory_percent(self):
        """Compare physical system memory to process resident memory and
        calculate process memory utilization as a percentage.
        """
        rss = _platform_impl.get_memory_info(self.pid)[0]
        try:
            return (rss / float(TOTAL_PHYMEM)) * 100
        except ZeroDivisionError:
            return 0.0

    def is_running(self):
        """Return whether the current process is running in the current process
        list."""
        try:
            new_proc = Process(self.pid)
            # calls get_process_info() which may in turn trigger NSP exception
            str(new_proc)
        except NoSuchProcess:
            return False
        return self == new_proc

    def kill(self, sig=None):
        """Kill the current process by using signal sig (defaults to SIGKILL).
        """
        # safety measure in case the current process has been killed in
        # meantime and the kernel reused its PID
        if not self.is_running():
            raise NoSuchProcess
        _platform_impl.kill_process(self.pid, sig)

def pid_exists(pid):
    """Check whether the given PID exists in the current process list."""
    return _platform_impl.pid_exists(pid)

def get_pid_list():
    """Return a list of current running PIDs."""
    return _platform_impl.get_pid_list()

def process_iter():
    """Return an iterator yielding a Process class instances for all
    running processes on the local machine.
    """
    pids = _platform_impl.get_pid_list()
    # for each PID, create a proxyied Process object
    # it will lazy init it's name and path later if required
    for pid in pids:
        try:
            yield Process(pid)
        except (NoSuchProcess, AccessDenied):
            continue

def get_process_list():
    """Return a list of Process class instances for all running
    processes on the local machine.
    """
    return list(process_iter())

def cpu_times():
    """Return system CPU times as a CPUTimes object."""
    values = get_system_cpu_times()
    return CPUTimes(**values)

_last_idle_time = cpu_times().idle
_last_time = time.time()

def cpu_percent():
    """Return the current system-wide CPU utilization as a percentage. For
    highest accuracy, it is recommended that this be called at least 1/10th
    of a second after importing the module or calling cpu_percent() in a
    previous call, to allow for a larger time delta from which to calculate
    the percentage value.
    """
    global _last_idle_time
    global _last_time

    idle_time = cpu_times().idle
    curr_time = time.time()

    time_delta = curr_time - _last_time
    idle_delta = idle_time - _last_idle_time

    # reset values for next run
    _last_idle_time = idle_time
    _last_time = curr_time

    # invalid data, will not be accurate so return 0.0 to avoid an overflow
    if time_delta < idle_delta:
        return 0.0

    try :
        idle_percent = (idle_delta / time_delta) * 100.0
        util_percent = (100 * NUM_CPUS) - idle_percent
    except ZeroDivisionError:
        return 0.0

    return util_percent / NUM_CPUS

def test():
    """List info of all currently running processes emulating a
    ps -aux output.
    """
    import datetime
    today_day = datetime.date.today()

    def get_process_info(pid):
        proc = Process(pid)
        uid = proc.uid
        pid = proc.pid
        cpu = round(proc.get_cpu_percent(), 1)
        mem = round(proc.get_memory_percent(), 1)
        rss, vsz = [x / 1024 for x in proc.get_memory_info()]

        # If process has been created today print H:M, else MonthDay
        start = datetime.datetime.fromtimestamp(proc.create_time)
        if start.date() == today_day:
            start = start.strftime("%H:%M")
        else:
            start = start.strftime("%b%d")

        cputime = time.strftime("%M:%S", time.localtime(sum(proc.get_cpu_times())))
        cmd = ' '.join(proc.cmdline)
        # where cmdline is not available UNIX shows process name between
        # [] parentheses
        if not cmd:
            cmd = "[%s]" %proc.name
        return "%-5s %7s %4s %4s %7s %7s %5s %8s %s" \
               %(uid, pid, cpu, mem, vsz, rss, start, cputime, cmd)

    print "%-5s %7s %4s %4s %7s %7s %5s %8s %s" \
       %("UID", "PID", "%CPU", "%MEM", "VSZ", "RSS", "START", "TIME", "COMMAND")
    pids = get_pid_list()
    pids.sort()
    for pid in pids:
        try:
            line = get_process_info(pid)
        except (AccessDenied, NoSuchProcess):
            pass
        else:
            print line

if __name__ == "__main__":
    test()
<commit_msg>(issue 59) modified is_running() to use only (pid + create time) for comparison resulting in a 10x speedup.<commit_after>#!/usr/bin/env python
#
# $Id$
#

"""psutil is a module providing convenience functions for managing
processes in a portable way by using Python.
"""

__all__ = [
    "NoSuchProcess",
    "AccessDenied",
    "NUM_CPUS",
    "ProcessInfo",
    "Process",
    "test",
    "CPUTimes",
    "pid_exists",
    "get_pid_list",
    "process_iter",
    "get_process_list",
    "TOTAL_PHYMEM",
    "avail_phymem",
    "used_phymem",
    "total_virtmem",
    "avail_virtmem",
    "used_virtmem",
    "cpu_times",
    "cpu_percent",
    ]

__version__ = '0.1.3'

import sys
import os
import time
try:
    import pwd, grp
except ImportError:
    pwd = grp = None

# exceptions are imported here, but may be overriden by platform
# module implementation later
from error import *

# import the appropriate module for our platform only
if sys.platform.lower().startswith("linux"):
    from _pslinux import *

elif sys.platform.lower().startswith("win32"):
    from _psmswindows import *

elif sys.platform.lower().startswith("darwin"):
    from _psosx import *

elif sys.platform.lower().startswith("freebsd"):
    from _psbsd import *
else:
    raise ImportError('no os specific module found')

# platform-specific modules define an Impl implementation class
_platform_impl = Impl()

class CPUTimes:
    """This class contains information about CPU times.
    It is not used directly but it's returned as an instance by
    psutil.cpu_times() function.

    Every CPU time is accessible in form of an attribute and represents
    the time CPU has spent in the given mode.

    The attributes availability varies depending on the platform.
    Here follows a list of all available attributes:

     - user
     - system
     - idle
     - nice (UNIX)
     - iowait (Linux)
     - irq (Linux, FreeBSD)
     - softirq (Linux)
    """

    def __init__(self, **kwargs):
        self.__attrs = []
        for name in kwargs:
            setattr(self, name, kwargs[name])
            self.__attrs.append(name)

    def __str__(self):
        string = []
        for attr in self.__attrs:
            value = getattr(self, attr)
            string.append("%s=%s" %(attr, value))
        return '; '.join(string)

    def __iter__(self):
        for attr in self.__attrs:
            yield getattr(self, attr)

class ProcessInfo(object):
    """Class that allows the process information to be passed between
    external code and psutil.  Used directly by the Process class.
    """

    def __init__(self, pid, ppid=None, name=None, path=None, cmdline=None,
                       uid=None, gid=None):
        self.pid = pid
        self.ppid = ppid
        self.name = name
        self.path = path
        self.cmdline = cmdline
        # if we have the cmdline but not the path, figure it out from argv[0]
        if cmdline and not path:
            self.path = os.path.dirname(cmdline[0])
        self.uid = uid
        self.gid = gid
        self.create = None
        self.username = None
        self.groupname = None

class Process(object):
    """Represents an OS process."""

    def __init__(self, pid):
        """Create a new Process object, raises NoSuchProcess if the PID
        does not exist, and ValueError if the parameter is not an
        integer PID."""
        if not isinstance(pid, int):
            raise ValueError("An integer is required")
        if not pid_exists(pid):
            raise NoSuchProcess("No process found with PID %s" % pid)
        self._procinfo = ProcessInfo(pid)
        self.is_proxy = True
        # try to init CPU times, if it raises AccessDenied then suppress
        # it so it won't interrupt the constructor. First call to
        # get_cpu_percent() will trigger the AccessDenied exception
        # instead.
        try:
            self._last_sys_time = time.time()
            self._last_user_time, self._last_kern_time = self.get_cpu_times()
        except AccessDenied:
            self._last_user_time, self._last_kern_time = None, None

    def __str__(self):
        return "psutil.Process <PID:%s; PPID:%s; NAME:'%s'; PATH:'%s'; " \
               "CMDLINE:%s; UID:%s; GID:%s;>" \
               %(self.pid, self.ppid, self.name, self.path, self.cmdline, \
                 self.uid, self.gid)

    def __eq__(self, other):
        """Test for equality with another Process object based on PID
        and creation time."""
        h1 = (self.pid, self.create_time)
        h2 = (other.pid, other.create_time)
        return h1 == h2

    def deproxy(self):
        """Used internally by Process properties. The first call to
        deproxy() initializes the ProcessInfo object in self._procinfo
        with process data read from platform-specific module's
        get_process_info() method.

        This method becomes a NO-OP after the first property is accessed.
        Property data is filled in from the ProcessInfo object created,
        and further calls to deproxy() simply return immediately without
        calling get_process_info().
        """
        if self.is_proxy:
            # get_process_info returns a tuple we use as the arguments
            # to the ProcessInfo constructor
            self._procinfo = ProcessInfo(*_platform_impl.get_process_info(self._procinfo.pid))
            self.is_proxy = False

    @property
    def pid(self):
        """The process pid."""
        return self._procinfo.pid

    @property
    def ppid(self):
        """The process parent pid."""
        self.deproxy()
        return self._procinfo.ppid

    @property
    def parent(self):
        """Return the parent process as a Process object. If no ppid is
        known then return None."""
        if self.ppid is not None:
            return Process(self.ppid)
        return None

    @property
    def name(self):
        """The process name."""
        self.deproxy()
        return self._procinfo.name

    @property
    def path(self):
        """The process path."""
        self.deproxy()
        return self._procinfo.path

    @property
    def cmdline(self):
        """The command line process has been called with."""
        self.deproxy()
        return self._procinfo.cmdline

    @property
    def uid(self):
        """The real user id of the current process."""
        self.deproxy()
        return self._procinfo.uid

    @property
    def gid(self):
        """The real group id of the current process."""
        self.deproxy()
        return self._procinfo.gid

    @property
    def username(self):
        """The name of the user that owns the process."""
    	if self._procinfo.username is not None:
    		return self._procinfo.username
        if pwd is not None:
            self._procinfo.username = pwd.getpwuid(self.uid).pw_name
        else:
            self._procinfo.username =  _platform_impl.get_process_username(self.pid)
    	return self._procinfo.username

    @property
    def groupname(self):
        """The real groupname of the current process."""
    	if self._procinfo.groupname is not None:
    		return self._procinfo.groupname
        if grp is not None:
            self._procinfo.groupname = grp.getgrgid(self.gid).gr_name
        else:
            # XXX - provisional for Windows
            return "groupname"
    	return self._procinfo.groupname

    @property
    def environ(self):
        """The process environment variables as a dictionary."""
        return _platform_impl.get_process_environ(self.pid)

    @property
    def create_time(self):
        """The process creation time as a floating point number
        expressed in seconds since the epoch, in UTC.
        """
        if self._procinfo.create is None:
            self._procinfo.create = _platform_impl.get_process_create_time(self.pid)
        return self._procinfo.create

    def getcwd(self):
        """Return a string representing the process current working directory."""
        return _platform_impl.get_process_cwd(self.pid)

    def get_cpu_percent(self):
        """Compare process times to system time elapsed since last call
        and calculate CPU utilization as a percentage. It is recommended
        for accuracy that this function be called with at least 1 second
        between calls. The initial delta is calculated from the
        instantiation of the Process object.
        """
        now = time.time()
        # will raise AccessDenied on OS X if not root or in procmod group
        user_t, kern_t = _platform_impl.get_cpu_times(self.pid)

        total_proc_time = float((user_t - self._last_user_time) + \
                                (kern_t - self._last_kern_time))
        try:
            percent = total_proc_time / float((now - self._last_sys_time))
        except ZeroDivisionError:
            percent = 0.000

        # reset the values
        self._last_sys_time = time.time()
        self._last_user_time, self._last_kern_time = self.get_cpu_times()

        return (percent * 100.0) / NUM_CPUS

    def get_cpu_times(self):
        """Return a tuple whose values are process CPU user and system
        time.  These are the same first two values that os.times()
        returns for the current process.
        """
        return _platform_impl.get_cpu_times(self.pid)

    def get_memory_info(self):
        """Return a tuple representing RSS (Resident Set Size) and VMS
        (Virtual Memory Size) in bytes.

        On UNIX RSS and VMS are the same values shown by ps.

        On Windows RSS and VMS refer to "Mem Usage" and "VM Size" columns
        of taskmgr.exe.
        """
        return _platform_impl.get_memory_info(self.pid)

    def get_memory_percent(self):
        """Compare physical system memory to process resident memory and
        calculate process memory utilization as a percentage.
        """
        rss = _platform_impl.get_memory_info(self.pid)[0]
        try:
            return (rss / float(TOTAL_PHYMEM)) * 100
        except ZeroDivisionError:
            return 0.0

    def is_running(self):
        """Return whether the current process is running in the current process
        list."""
        try:
            newproc = Process(self.pid)
            return self == newproc
        except NoSuchProcess:
            return False

    def kill(self, sig=None):
        """Kill the current process by using signal sig (defaults to SIGKILL).
        """
        # safety measure in case the current process has been killed in
        # meantime and the kernel reused its PID
        if not self.is_running():
            raise NoSuchProcess
        _platform_impl.kill_process(self.pid, sig)

def pid_exists(pid):
    """Check whether the given PID exists in the current process list."""
    return _platform_impl.pid_exists(pid)

def get_pid_list():
    """Return a list of current running PIDs."""
    return _platform_impl.get_pid_list()

def process_iter():
    """Return an iterator yielding a Process class instances for all
    running processes on the local machine.
    """
    pids = _platform_impl.get_pid_list()
    # for each PID, create a proxyied Process object
    # it will lazy init it's name and path later if required
    for pid in pids:
        try:
            yield Process(pid)
        except (NoSuchProcess, AccessDenied):
            continue

def get_process_list():
    """Return a list of Process class instances for all running
    processes on the local machine.
    """
    return list(process_iter())

def cpu_times():
    """Return system CPU times as a CPUTimes object."""
    values = get_system_cpu_times()
    return CPUTimes(**values)

_last_idle_time = cpu_times().idle
_last_time = time.time()

def cpu_percent():
    """Return the current system-wide CPU utilization as a percentage. For
    highest accuracy, it is recommended that this be called at least 1/10th
    of a second after importing the module or calling cpu_percent() in a
    previous call, to allow for a larger time delta from which to calculate
    the percentage value.
    """
    global _last_idle_time
    global _last_time

    idle_time = cpu_times().idle
    curr_time = time.time()

    time_delta = curr_time - _last_time
    idle_delta = idle_time - _last_idle_time

    # reset values for next run
    _last_idle_time = idle_time
    _last_time = curr_time

    # invalid data, will not be accurate so return 0.0 to avoid an overflow
    if time_delta < idle_delta:
        return 0.0

    try :
        idle_percent = (idle_delta / time_delta) * 100.0
        util_percent = (100 * NUM_CPUS) - idle_percent
    except ZeroDivisionError:
        return 0.0

    return util_percent / NUM_CPUS

def test():
    """List info of all currently running processes emulating a
    ps -aux output.
    """
    import datetime
    today_day = datetime.date.today()

    def get_process_info(pid):
        proc = Process(pid)
        uid = proc.uid
        pid = proc.pid
        cpu = round(proc.get_cpu_percent(), 1)
        mem = round(proc.get_memory_percent(), 1)
        rss, vsz = [x / 1024 for x in proc.get_memory_info()]

        # If process has been created today print H:M, else MonthDay
        start = datetime.datetime.fromtimestamp(proc.create_time)
        if start.date() == today_day:
            start = start.strftime("%H:%M")
        else:
            start = start.strftime("%b%d")

        cputime = time.strftime("%M:%S", time.localtime(sum(proc.get_cpu_times())))
        cmd = ' '.join(proc.cmdline)
        # where cmdline is not available UNIX shows process name between
        # [] parentheses
        if not cmd:
            cmd = "[%s]" %proc.name
        return "%-5s %7s %4s %4s %7s %7s %5s %8s %s" \
               %(uid, pid, cpu, mem, vsz, rss, start, cputime, cmd)

    print "%-5s %7s %4s %4s %7s %7s %5s %8s %s" \
       %("UID", "PID", "%CPU", "%MEM", "VSZ", "RSS", "START", "TIME", "COMMAND")
    pids = get_pid_list()
    pids.sort()
    for pid in pids:
        try:
            line = get_process_info(pid)
        except (AccessDenied, NoSuchProcess):
            pass
        else:
            print line

if __name__ == "__main__":
    test()
<|endoftext|>


<commit_before>#!/usr/bin/env python

"""
BenchIT is a simple Python script for security auditing purposes.

It is the most useful for auditing systems by analyzing files collected from
remote hosts. It allows you to check certain system settings stored in
configuration files by using regular expressions. The simplest check will
determine if a pattern matches. However, with capturing groups it is also
possible to determine whether the actual value matches the expected or the
default value. Results are summarized in a HTML report and also stored in a
CSV file.

Regular expressions must be stored in a CSV file containing the following
columns separated by a semicolon (;):
  - Boolean value which determines whether we expect a match or not ("True").
  - Name of the file ("sshd_config").
  - Regular expression to evaluate ("^Protocol[ \t]+(\d)$").
  - Number of the chaper in the given CIS guideline ("9.2.1").
  - Title of the chapter in the given CIS guideline ("Set SSH Protocol to 2").
  - Short summary of the chapter ("SSH supports two different protocols...").
  - Default value of the setting ("2,1").
  - Expected value of the setting ("2").

system requirements:
  - GnuWin32 with ports of the necessary tools
        http://getgnuwin32.sourceforge.net/
  - Q to execute SQL-like queries on CSV files
        http://harelba.github.io/q/
  - The following Python libraries:
        dominate, configobj
"""

__description__ = 'Simple Python script for security auditing purposes.'
__author__ = 'Gabor Seljan'
__version__ = '0.1.6'
__date__ = '2016/11/12'

import io
import os
import re
import csv
import ast
import sys
import time
import binascii
import textwrap
import dominate

from argparse import *
from dominate.tags import *
from configobj import ConfigObj
from subprocess import check_output

banner = """
                     ____                  _     _____ _______
                    |  _ \     v{}     | |   |_   _|__   __|
                    | |_) | ___ _ __   ___| |__   | |    | |
                    |  _ < / _ \ '_ \ / __| '_ \  | |    | |
                    | |_) |  __/ | | | (__| | | |_| |_   | |
                    |____/ \___|_| |_|\___|_| |_|_____|  |_|
""".format(__version__)

print(banner)

parser = ArgumentParser(
    formatter_class=RawDescriptionHelpFormatter,
    description=__doc__,
    prog='benchit'
)

module = parser.add_mutually_exclusive_group(required=True)
module.add_argument('-d, --database', dest='d', action='store_true',
                    help='audit Oracle database')
module.add_argument('-l, --linux', dest='l', action='store_true',
                    help='audit Linux system')
module.add_argument('-w, --windows', dest='w', action='store_true',
                    help='audit Windows system')

parser.add_argument('-o, --output', dest='o', default='results',
                    help='output filename (default results_{timestamp}.html)')
parser.add_argument('-p, --path', dest='p', default='.',
                    help='base path to target directory (default .)')
parser.add_argument('-v, --verbose', dest='v', action='store_true',
                    help='run in verbose mode')
parser.add_argument('-s, --skip-dirlist', dest='s', action='store_true',
                    help='skip directory list checking (default false)')
parser.add_argument('--debug', action='store_true',
                    help='run in debug mode (default false)')

args = parser.parse_args()

items = {}
audit = {}
results = []
total = 0
passed = 0
failed = 0
errors = 0

headers = [
    'Chapter',
    'Title',
    'Summary',
    'Default',
    'Actual',
    'Expected',
    'Result'
]

def main():
    global audit, results, total, passed, failed, errors

    config = ConfigObj('benchit.ini')

    if args.d:
        print('[+] Oracle detected!')
        audit = config['Database']['Oracle']
    elif args.w:
        print('[+] Windows detected!')
        audit = config['Windows']['2012']
    elif args.l:
        for l in ['RedHat', 'SuSe', 'LSB', 'Debian']:
            if os.path.isfile('%s/etc/%s-release' % (args.p, l)):
                print('[+] %s Linux detected!' % l)
                audit = config['Linux'][l]
                break

    if audit is None:
        print('[!] Nothing to audit...')
        exit(1)

    with open(audit['csv'], mode='r') as infile:
        reader = csv.reader(infile, delimiter=';')
        for row in reader:
            i = row[0]
            n = row[1]
            if i in items:
                if n in items[i]:
                    items[i][n].append(tuple(row[2:]))
                else:
                    items[i][n] = [tuple(row[2:])]
            else:
                items[i] = {row[1]: [tuple(row[2:])]}

    print('[*] Audit in progress, this can take a while...')

    for r, d in items.items():
        print('[*]   Checking %s items...' % r)
        if r != 'None':
            for f, i in d.items():
                print('[+]     Processing %s' % f)
                path = '/'.join([args.p, f])
                if args.d:
                    check_item_database(path, i, ast.literal_eval(r))
                elif args.l or args.w:
                    check_item_os(path, i, ast.literal_eval(r))
        else:
            for f, i in d.items():
                try:
                    for c, n, t, d, *_ in i:
                        total += 1
                        if not os.path.isfile('/'.join([args.p, f])):
                            raise IOError('{} not found!'.format(f))
                        c = c.format(args.p)
                        print('[+]     Executing {}'.format(c))
                        output = check_output(c, shell=True)
                        if output:
                            failed += 1
                            r = 'Fail'
                        else:
                            passed += 1
                            r = 'Pass'
                        results.append((n, t, d, '', '', '', r))
                except (IOError, OSError) as err:
                    print('[!] Error: {}'.format(err))
                    for c, n, t, d, *_ in i:
                        errors += 1
                        results.append((n, t, d, '', '', '', 'Error'))

    timestamp = time.strftime("%Y%m%dT%H%M%S")
    filename = '{}_{}'.format(args.o, timestamp)
    create_html_report('{}.html'.format(filename))
    create_csv_report('{}.csv'.format(filename))

    print('[*] Audit finished!')

def check_item_os(filename, items, expected):
    """Checks every regex pattern listed in the loaded CSV file."""
    global results, total, passed, failed, errors
    # Convert UTF-16 encoded files (created by reg export) to UTF-8
    if args.w:
        try:
            with io.open(filename, 'r', encoding='utf-16') as f:
                text = f.read()
            with io.open(filename, 'w', encoding='utf-8') as f:
                f.write(text)
        except (UnicodeError, UnicodeDecodeError, IOError) as e:
            pass
    if args.l and args.s and ('dirlist.txt' in filename):
        return
    try:
        with open(filename, "r") as f:
            s = f.read().replace('\\', '\\\\')
            for p, n, t, b, d, e in items:
                total += 1
                m = re.search(p, s, re.M)
                e = 'N/A' if len(e) == 0 else e
                m = 'N/A' if m is None or len(m.groups()) == 0 else m.group(1)
                # Convert null-terminated strings from HEX to ASCII
                if args.w and m.startswith('hex'):
                    m = re.sub('(00,?|[\s,]|\\\\)', '', m[7:])
                    m = binascii.unhexlify(m)
                    m = str(m)[2:-1]
                # We expect a match
                if expected:
                    # We found a match
                    if m != 'N/A':
                        # Any value is accepted
                        if e == 'N/A':
                            passed += 1
                            r = 'Pass'
                        # Check relational (>, <, =) values
                        elif check_item_relational(m, e):
                            passed += 1
                            r = 'Pass'
                        else:
                            failed += 1
                            r = 'Fail'
                    # We have a default value
                    elif len(d) != 0 and check_item_default(d, e):
                            passed += 1
                            r = 'Pass'
                    else:
                        failed += 1
                        r = 'Fail'
                # We don't expect a match
                elif m == 'N/A' and not expected:
                    passed += 1
                    r = 'Pass'
                else:
                    failed += 1
                    r = 'Fail'
                results.append((n, t, b, d, m, e, r))
                if args.debug:
                    msg = 'N:{}\tD:{}\tM:{}\tE:{}\tR:{}'.format(n, d, m, e, r)
                    print('        {}'.format(msg))
    except IOError as err:
        print('[!] Error: %s not found!' % filename)
        if err.errno is 2:
            for p, n, t, b, d, e in items:
                total += 1
                m = 'N/F'
                if len(d) != 0:
                    if check_item_default(d, e):
                        passed += 1
                        r = 'Pass'
                    else:
                        failed += 1
                        r = 'Fail'
                else:
                    errors += 1
                    r = 'Error'
                results.append((n, t, b, d, m, e, r))
                if args.debug:
                    msg = 'N:{}\tD:{}\tM:{}\tE:{}\tR:{}'.format(n, d, m, e, r)
                    print('        {}'.format(msg))
        return

def check_item_default(d, e):
    """Checks the default value when the specific setting is not available."""
    if e[:1] in ['>', '<'] and d not in ['Not defined', 'Not configured']:
        return check_item_relational(d, e)
    elif d == e:
        return True
    else:
        return False

def check_item_relational(a, e):
    """Checks for acceptable lesser or greather values."""
    if e[:1] == '>' and int(a) >= int(e[1:]):
        return True
    elif e[:1] == '<' and int(a) <= int(e[1:]):
        return True
    elif a == e:
        return True
    else:
        return False

def check_item_database(filename, items, expected):
    """Checks every query listed in the loaded CSV file."""
    global results, total, passed, failed, errors
    try:
        for q, n, t, b, d, e in items:
            if args.v:
                print('[+]            {}'.format(q))
            if not os.path.isfile(filename):
                raise IOError('{} not found!'.format(filename))
            total += 1
            params = ['q', '-H', '-d', ';', q.format(filename)]
            o = str(check_output(params, shell=True).strip())
            if o.startswith('b\''):
                o = o[2:-1]
            if o and e:
                if args.debug:
                    print(
                        'o: {:s} {:s} e: {:s} {:s}'.format(
                            o.strip(),
                            type(o.strip()),
                            e.strip(),
                            type(e.strip())
                            )
                        )
                if o != e:
                    failed += 1
                    results.append((n, t, b, d, o, e, 'Fail'))
                else:
                    passed += 1
                    results.append((n, t, b, d, o, e, 'Pass'))
            elif bool(o) is not expected:
                failed += 1
                results.append((n, t, b, d, '', '', 'Fail'))
            else:
                passed += 1
                results.append((n, t, b, d, '', '', 'Pass'))
    except (IOError, OSError) as err:
        print('[!] Error: {}'.format(err))
        for q, n, t, b, d, e in items:
            total += 1
            errors += 1
            results.append((n, t, b, d, '', e, 'Error'))
        return

def create_html_report(filename):
    """Creates a nice HTML report from the results."""
    title = 'Security Audit Report ({})'.format(time.strftime('%m/%d/%Y'))
    doc = dominate.document(title=title)
    with doc:
        styles = [
            'margin:0;padding:0;',
            'border:1px solid black;',
            'margin-top:0;',
            'margin-bottom:0',
        ]
        colors = {
            'Pass': 'lime',
            'Fail': 'red',
            'Error': 'yellow',
        }

        for line in banner.splitlines():
            line = re.sub(r' {20}', '     ', line)
            pre(line, style='margin:0;font-size:16px;font-weight:bold')
        h1(title, style=styles[3])
        h4(audit['benchmark'], style=styles[0])

        p('Performed %d tests in total:' % total, style=styles[3])
        l = ul(style=styles[2])
        l += li('Pass = {} ({:.0f}%)'.format(passed, passed/float(total)*100))
        l += li('Fail = {} ({:.0f}%)'.format(failed, failed/float(total)*100))
        l += li('Error = {} ({:.0f}%)'.format(errors, errors/float(total)*100))

        t = table(
            border=1,
            width='100%',
            cellspacing='0',
            cellpadding='3',
            style=styles[1]
        )

        with t.add(thead(border=1, style=styles[1])):
            l = tr()
            for header in headers:
                l += th(header, bgcolor='black', style='color:white')
        with t.add(tbody(border=1, style=styles[1])):
            for n, t, b, d, a, e, r in sorted(set(results), key=lambda x: tuple(map(int, x[0].split('.')))):
                l = tr(style='border:1px solid black')
                l += td(n, align='left', width='5%', style=styles[1])
                l += td(t, align='left', width='32%', style=styles[1])
                l += td(b, align='left', width='30%', style=styles[1])
                l += td(d, align='left', width='10%', style=styles[1])
                l += td(a, align='left', width='10%', style=styles[1])
                l += td(e, align='left', width='10%', style=styles[1])
                l += td(r, align='center', bgcolor=colors[r], style=styles[1])

    with open(filename, 'w') as f:
        f.write(str(doc))
        print('[+]   Created %s' % filename)

def create_csv_report(filename):
    """Creates a CSV report from the results."""
    with open(filename, 'w', newline='') as f:
        w = csv.writer(f, delimiter=';')
        w.writerow(headers)
        for n, t, b, d, a, e, r in sorted(set(results), key=lambda x: tuple(map(int, x[0].split('.')))):
            w.writerow([n, t, b, d, a, e, r])
        print('[+]   Created %s' % filename)

if __name__ == "__main__":
        main()
<commit_msg>Explicitly specify the encoding of the file<commit_after>#!/usr/bin/env python

"""
BenchIT is a simple Python script for security auditing purposes.

It is the most useful for auditing systems by analyzing files collected from
remote hosts. It allows you to check certain system settings stored in
configuration files by using regular expressions. The simplest check will
determine if a pattern matches. However, with capturing groups it is also
possible to determine whether the actual value matches the expected or the
default value. Results are summarized in a HTML report and also stored in a
CSV file.

Regular expressions must be stored in a CSV file containing the following
columns separated by a semicolon (;):
  - Boolean value which determines whether we expect a match or not ("True").
  - Name of the file ("sshd_config").
  - Regular expression to evaluate ("^Protocol[ \t]+(\d)$").
  - Number of the chaper in the given CIS guideline ("9.2.1").
  - Title of the chapter in the given CIS guideline ("Set SSH Protocol to 2").
  - Short summary of the chapter ("SSH supports two different protocols...").
  - Default value of the setting ("2,1").
  - Expected value of the setting ("2").

system requirements:
  - GnuWin32 with ports of the necessary tools
        http://getgnuwin32.sourceforge.net/
  - Q to execute SQL-like queries on CSV files
        http://harelba.github.io/q/
  - The following Python libraries:
        dominate, configobj
"""

__description__ = 'Simple Python script for security auditing purposes.'
__author__ = 'Gabor Seljan'
__version__ = '0.1.7'
__date__ = '2016/11/26'

import io
import os
import re
import csv
import ast
import sys
import time
import binascii
import textwrap
import dominate

from argparse import *
from dominate.tags import *
from configobj import ConfigObj
from subprocess import check_output

banner = """
                     ____                  _     _____ _______
                    |  _ \     v{}     | |   |_   _|__   __|
                    | |_) | ___ _ __   ___| |__   | |    | |
                    |  _ < / _ \ '_ \ / __| '_ \  | |    | |
                    | |_) |  __/ | | | (__| | | |_| |_   | |
                    |____/ \___|_| |_|\___|_| |_|_____|  |_|
""".format(__version__)

print(banner)

parser = ArgumentParser(
    formatter_class=RawDescriptionHelpFormatter,
    description=__doc__,
    prog='benchit'
)

module = parser.add_mutually_exclusive_group(required=True)
module.add_argument('-d, --database', dest='d', action='store_true',
                    help='audit Oracle database')
module.add_argument('-l, --linux', dest='l', action='store_true',
                    help='audit Linux system')
module.add_argument('-w, --windows', dest='w', action='store_true',
                    help='audit Windows system')

parser.add_argument('-o, --output', dest='o', default='results',
                    help='output filename (default results_{timestamp}.html)')
parser.add_argument('-p, --path', dest='p', default='.',
                    help='base path to target directory (default .)')
parser.add_argument('-v, --verbose', dest='v', action='store_true',
                    help='run in verbose mode')
parser.add_argument('-s, --skip-dirlist', dest='s', action='store_true',
                    help='skip directory list checking (default false)')
parser.add_argument('--debug', action='store_true',
                    help='run in debug mode (default false)')

args = parser.parse_args()

items = {}
audit = {}
results = []
total = 0
passed = 0
failed = 0
errors = 0

headers = [
    'Chapter',
    'Title',
    'Summary',
    'Default',
    'Actual',
    'Expected',
    'Result'
]

def main():
    global audit, results, total, passed, failed, errors

    config = ConfigObj('benchit.ini')

    if args.d:
        print('[+] Oracle detected!')
        audit = config['Database']['Oracle']
    elif args.w:
        print('[+] Windows detected!')
        audit = config['Windows']['2012']
    elif args.l:
        for l in ['RedHat', 'SuSe', 'LSB', 'Debian']:
            if os.path.isfile('%s/etc/%s-release' % (args.p, l)):
                print('[+] %s Linux detected!' % l)
                audit = config['Linux'][l]
                break

    if audit is None:
        print('[!] Nothing to audit...')
        exit(1)

    with open(audit['csv'], mode='r') as infile:
        reader = csv.reader(infile, delimiter=';')
        for row in reader:
            i = row[0]
            n = row[1]
            if i in items:
                if n in items[i]:
                    items[i][n].append(tuple(row[2:]))
                else:
                    items[i][n] = [tuple(row[2:])]
            else:
                items[i] = {row[1]: [tuple(row[2:])]}

    print('[*] Audit in progress, this can take a while...')

    for r, d in items.items():
        print('[*]   Checking %s items...' % r)
        if r != 'None':
            for f, i in d.items():
                print('[+]     Processing %s' % f)
                path = '/'.join([args.p, f])
                if args.d:
                    check_item_database(path, i, ast.literal_eval(r))
                elif args.l or args.w:
                    check_item_os(path, i, ast.literal_eval(r))
        else:
            for f, i in d.items():
                try:
                    for c, n, t, d, *_ in i:
                        total += 1
                        if not os.path.isfile('/'.join([args.p, f])):
                            raise IOError('{} not found!'.format(f))
                        c = c.format(args.p)
                        print('[+]     Executing {}'.format(c))
                        output = check_output(c, shell=True)
                        if output:
                            failed += 1
                            r = 'Fail'
                        else:
                            passed += 1
                            r = 'Pass'
                        results.append((n, t, d, '', '', '', r))
                except (IOError, OSError) as err:
                    print('[!] Error: {}'.format(err))
                    for c, n, t, d, *_ in i:
                        errors += 1
                        results.append((n, t, d, '', '', '', 'Error'))

    timestamp = time.strftime("%Y%m%dT%H%M%S")
    filename = '{}_{}'.format(args.o, timestamp)
    create_html_report('{}.html'.format(filename))
    create_csv_report('{}.csv'.format(filename))

    print('[*] Audit finished!')

def check_item_os(filename, items, expected):
    """Checks every regex pattern listed in the loaded CSV file."""
    global results, total, passed, failed, errors
    # Convert UTF-16 encoded files (created by reg export) to UTF-8
    if args.w:
        try:
            with io.open(filename, 'r', encoding='utf-16') as f:
                text = f.read()
            with io.open(filename, 'w', encoding='utf-8') as f:
                f.write(text)
        except (UnicodeError, UnicodeDecodeError, IOError) as e:
            pass
    if args.l and args.s and ('dirlist.txt' in filename):
        return
    try:
        with open(filename, 'r', encoding='utf8') as f:
            s = f.read().replace('\\', '\\\\')
            for p, n, t, b, d, e in items:
                total += 1
                m = re.search(p, s, re.M)
                e = 'N/A' if len(e) == 0 else e
                m = 'N/A' if m is None or len(m.groups()) == 0 else m.group(1)
                # Convert null-terminated strings from HEX to ASCII
                if args.w and m.startswith('hex'):
                    m = re.sub('(00,?|[\s,]|\\\\)', '', m[7:])
                    m = binascii.unhexlify(m)
                    m = str(m)[2:-1]
                # We expect a match
                if expected:
                    # We found a match
                    if m != 'N/A':
                        # Any value is accepted
                        if e == 'N/A':
                            passed += 1
                            r = 'Pass'
                        # Check relational (>, <, =) values
                        elif check_item_relational(m, e):
                            passed += 1
                            r = 'Pass'
                        else:
                            failed += 1
                            r = 'Fail'
                    # We have a default value
                    elif len(d) != 0 and check_item_default(d, e):
                            passed += 1
                            r = 'Pass'
                    else:
                        failed += 1
                        r = 'Fail'
                # We don't expect a match
                elif m == 'N/A' and not expected:
                    passed += 1
                    r = 'Pass'
                else:
                    failed += 1
                    r = 'Fail'
                results.append((n, t, b, d, m, e, r))
                if args.debug:
                    msg = 'N:{}\tD:{}\tM:{}\tE:{}\tR:{}'.format(n, d, m, e, r)
                    print('        {}'.format(msg))
    except IOError as err:
        print('[!] Error: %s not found!' % filename)
        if err.errno is 2:
            for p, n, t, b, d, e in items:
                total += 1
                m = 'N/F'
                if len(d) != 0:
                    if check_item_default(d, e):
                        passed += 1
                        r = 'Pass'
                    else:
                        failed += 1
                        r = 'Fail'
                else:
                    errors += 1
                    r = 'Error'
                results.append((n, t, b, d, m, e, r))
                if args.debug:
                    msg = 'N:{}\tD:{}\tM:{}\tE:{}\tR:{}'.format(n, d, m, e, r)
                    print('        {}'.format(msg))
        return

def check_item_default(d, e):
    """Checks the default value when the specific setting is not available."""
    if e[:1] in ['>', '<'] and d not in ['Not defined', 'Not configured']:
        return check_item_relational(d, e)
    elif d == e:
        return True
    else:
        return False

def check_item_relational(a, e):
    """Checks for acceptable lesser or greather values."""
    if e[:1] == '>' and int(a) >= int(e[1:]):
        return True
    elif e[:1] == '<' and int(a) <= int(e[1:]):
        return True
    elif a == e:
        return True
    else:
        return False

def check_item_database(filename, items, expected):
    """Checks every query listed in the loaded CSV file."""
    global results, total, passed, failed, errors
    try:
        for q, n, t, b, d, e in items:
            if args.v:
                print('[+]            {}'.format(q))
            if not os.path.isfile(filename):
                raise IOError('{} not found!'.format(filename))
            total += 1
            params = ['q', '-H', '-d', ';', q.format(filename)]
            o = str(check_output(params, shell=True).strip())
            if o.startswith('b\''):
                o = o[2:-1]
            if o and e:
                if args.debug:
                    print(
                        'o: {:s} {:s} e: {:s} {:s}'.format(
                            o.strip(),
                            type(o.strip()),
                            e.strip(),
                            type(e.strip())
                            )
                        )
                if o != e:
                    failed += 1
                    results.append((n, t, b, d, o, e, 'Fail'))
                else:
                    passed += 1
                    results.append((n, t, b, d, o, e, 'Pass'))
            elif bool(o) is not expected:
                failed += 1
                results.append((n, t, b, d, '', '', 'Fail'))
            else:
                passed += 1
                results.append((n, t, b, d, '', '', 'Pass'))
    except (IOError, OSError) as err:
        print('[!] Error: {}'.format(err))
        for q, n, t, b, d, e in items:
            total += 1
            errors += 1
            results.append((n, t, b, d, '', e, 'Error'))
        return

def create_html_report(filename):
    """Creates a nice HTML report from the results."""
    title = 'Security Audit Report ({})'.format(time.strftime('%m/%d/%Y'))
    doc = dominate.document(title=title)
    with doc:
        styles = [
            'margin:0;padding:0;',
            'border:1px solid black;',
            'margin-top:0;',
            'margin-bottom:0',
        ]
        colors = {
            'Pass': 'lime',
            'Fail': 'red',
            'Error': 'yellow',
        }

        for line in banner.splitlines():
            line = re.sub(r' {20}', '     ', line)
            pre(line, style='margin:0;font-size:16px;font-weight:bold')
        h1(title, style=styles[3])
        h4(audit['benchmark'], style=styles[0])

        p('Performed %d tests in total:' % total, style=styles[3])
        l = ul(style=styles[2])
        l += li('Pass = {} ({:.0f}%)'.format(passed, passed/float(total)*100))
        l += li('Fail = {} ({:.0f}%)'.format(failed, failed/float(total)*100))
        l += li('Error = {} ({:.0f}%)'.format(errors, errors/float(total)*100))

        t = table(
            border=1,
            width='100%',
            cellspacing='0',
            cellpadding='3',
            style=styles[1]
        )

        with t.add(thead(border=1, style=styles[1])):
            l = tr()
            for header in headers:
                l += th(header, bgcolor='black', style='color:white')
        with t.add(tbody(border=1, style=styles[1])):
            for n, t, b, d, a, e, r in sorted(set(results), key=lambda x: tuple(map(int, x[0].split('.')))):
                l = tr(style='border:1px solid black')
                l += td(n, align='left', width='5%', style=styles[1])
                l += td(t, align='left', width='32%', style=styles[1])
                l += td(b, align='left', width='30%', style=styles[1])
                l += td(d, align='left', width='10%', style=styles[1])
                l += td(a, align='left', width='10%', style=styles[1])
                l += td(e, align='left', width='10%', style=styles[1])
                l += td(r, align='center', bgcolor=colors[r], style=styles[1])

    with open(filename, 'w') as f:
        f.write(str(doc))
        print('[+]   Created %s' % filename)

def create_csv_report(filename):
    """Creates a CSV report from the results."""
    with open(filename, 'w', newline='') as f:
        w = csv.writer(f, delimiter=';')
        w.writerow(headers)
        for n, t, b, d, a, e, r in sorted(set(results), key=lambda x: tuple(map(int, x[0].split('.')))):
            w.writerow([n, t, b, d, a, e, r])
        print('[+]   Created %s' % filename)

if __name__ == "__main__":
        main()
<|endoftext|>


<commit_before>"""Base models for "Content", including the indexing and search features
that we want any piece of content to have."""

from django.conf import settings
from django.contrib.contenttypes.models import ContentType
from django.core.urlresolvers import reverse
from django.db import models
from django.db.backends import util
from django.template.defaultfilters import slugify
from django.utils import timezone

from bulbs.images.fields import RemoteImageField

from elasticutils import SearchResults, S
from elasticutils.contrib.django import get_es
from polymorphic import PolymorphicModel

class ReadonlyRelatedManager(object):
    """Replaces Django's RelatedMangers in read-only scenarios."""
    def __init__(self):
        self.data = []

    def __set__(self, obj, data):
        if data:
            if not isinstance(data, list):
                data = [data]
        else:
            data = []
        self.data = data

    def add(self, *args):
        self._on_bad_access()

    def all(self):
        return self.data

    def clear(self):
        self._on_bad_access()

    def remove(self, *args):
        self._on_bad_access()

    def _on_bad_access(self):
        raise TypeError('%s is read-only' % self.__class__.__name__)

def readonly_content_factory(model):
    """This is used to generate a class for a piece of content that we
    have retrieved from ElasticSearch. This content will have a number
    of "read-only" fields."""

    class Meta:
        proxy = True
        app_label = model._meta.app_label

    name = '%s_Readonly' % model.__name__
    name = util.truncate_name(name, 80, 32)

    overrides = {
        'Meta': Meta,
        '__module__': model.__module__,
        '_readonly': True,
        'tags': ReadonlyRelatedManager()
    }
    # TODO: Add additional deferred fields here, just as placeholders.
    return type(str(name), (model,), overrides)

def deserialize_polymorphic_model(data):
    """Deserializes simple polymorphic models."""
    content_type = ContentType.objects.get_for_id(data['polymorphic_ctype_id'])
    if content_type:
        klass = content_type.model_class()
        instance = klass.from_source(data)
        return instance

class ContentSearchResults(SearchResults):
    def set_objects(self, results):
        self.objects = []
        readonly_cls = readonly_content_factory(Content)
        for result in results:
            obj = readonly_cls.from_source(result['_source'])
            self.objects.append(obj)

    def __iter__(self):
        return self.objects.__iter__()

class ContentS(S):

    def all(self):
        """
        Fixes the default `S.all` method given by elasticutils.
        `S` generally looks like django queryset but differs in
        a few ways, one of which is `all`. Django `QuerySet` just
        returns a clone for `all` but `S` wants to return all
        the documents. This makes `all` at least respect slices
        but the real fix is to probably make `S` work more like
        `QuerySet`.
        """
        if self.start == 0 and self.stop is None:
            # no slicing has occurred. let's get all of the records.
            count = self.count()
            return self[:count].execute()
        return self.execute()

    def get_results_class(self):
        """Returns the results class to use

        The results class should be a subclass of SearchResults.

        """
        if self.as_list or self.as_dict:
            return super(ContentS, self).get_results_class()

        return ContentSearchResults

class TagSearchResults(SearchResults):
    def set_objects(self, results):
        self.objects = [
            deserialize_polymorphic_model(result['_source']) for result in results
        ]

    def __iter__(self):
        return self.objects.__iter__()

class TagS(S):

    def get_results_class(self):
        """Returns the results class to use

        The results class should be a subclass of SearchResults.

        """
        if self.as_list or self.as_dict:
            return super(TagS, self).get_results_class()

        return TagSearchResults

class PolymorphicIndexable(object):
    """Base mixin for polymorphic indexin'"""
    def extract_document(self):
        return {
            'id': self.id,
            'polymorphic_ctype_id': self.polymorphic_ctype_id
        }

    def index(self, refresh=False):
        es = get_es(urls=settings.ES_URLS)
        index = settings.ES_INDEXES.get('default')
        es.index(
            index,
            self.get_mapping_type_name(),
            self.extract_document(),
            self.id,
            refresh=refresh
        )

    def save(self, index=True, refresh=False, *args, **kwargs):
        result = super(PolymorphicIndexable, self).save(*args, **kwargs)
        if index:
            self.index(refresh=refresh)
        return result

    @classmethod
    def get_mapping(cls):
        return {
            cls.get_mapping_type_name(): {
                'properties': cls.get_mapping_properties()
            }
        }

    @classmethod
    def get_mapping_properties(cls):
        return {
            'id': {'type': 'integer'},
            'polymorphic_ctype_id': {'type': 'integer'}
        }

    @classmethod
    def get_mapping_type_name(cls):
        return '%s_%s' % (cls._meta.app_label, cls.__name__.lower())

class Tag(PolymorphicIndexable, PolymorphicModel):
    """Model for tagging up Content."""
    name = models.CharField(max_length=255)
    slug = models.SlugField()

    _doctype_cache = {}

    def __unicode__(self):
        return '%s: %s' % (self.__class__.__name__, self.name)

    def save(self, *args, **kwargs):
        self.slug = slugify(self.name)
        return super(Tag, self).save(*args, **kwargs)

    def extract_document(self):
        doc = super(Tag, self).extract_document()
        doc.update({
            'name': self.name,
            'slug': self.slug
        })
        return doc

    @classmethod
    def from_source(cls, _source):
        return cls(
            id=_source['id'],
            polymorphic_ctype_id=_source['polymorphic_ctype_id'],
            name=_source['name'],
            slug=_source['slug']
        )

    @classmethod
    def get_doctypes(cls):
        if len(cls._doctype_cache) == 0:
            for app in models.get_apps():
                for model in models.get_models(app, include_auto_created=True):
                    if isinstance(model(), Tag):
                        cls._doctype_cache[model.get_mapping_type_name()] = model
        return cls._doctype_cache

    @classmethod
    def get_mapping_properties(cls):
        props = super(Tag, cls).get_mapping_properties()
        props.update({
            'name': {'type': 'string', 'index': 'not_analyzed'},
            'slug': {'type': 'string', 'index': 'not_analyzed'},
        })
        return props

    @classmethod
    def search(cls, **kwargs):
        """Search tags...profit."""
        index = settings.ES_INDEXES.get('default')
        results = TagS().es(urls=settings.ES_URLS).indexes(index)
        name = kwargs.pop('name', '')
        if name:
            results = results.query(name__prefix=name, boost=4, should=True).query(name__fuzzy={
                'value': name,
                'prefix_length': 1,
                'min_similarity': 0.35
            }, should=True)

        types = kwargs.pop('types', [])
        if types:
            # only use valid subtypes
            results = results.doctypes(*[
                type_classname for type_classname in kwargs['types'] \
                if type_classname in cls.get_doctypes()
            ])
        else:
            results = results.doctypes(*cls.get_doctypes().keys())
        return results

class Section(Tag):
    """Tag subclass which represents major sections of the site."""
    class Meta(Tag.Meta):
        proxy = True

class Content(PolymorphicIndexable, PolymorphicModel):
    """The base content model from which all other content derives."""
    published = models.DateTimeField(blank=True, null=True)
    title = models.CharField(max_length=512)
    slug = models.SlugField()
    description = models.TextField(max_length=1024)
    image = RemoteImageField(null=True, blank=True)

    authors = models.ManyToManyField(settings.AUTH_USER_MODEL)
    _byline = models.CharField(max_length=255, null=True, blank=True)  # This is an overridable field that is by default the author names
    _tags = models.TextField(null=True, blank=True)  # A return-separated list of tag names, exposed as a list of strings
    _feature_type = models.CharField(max_length=255, null=True, blank=True)  # "New in Brief", "Newswire", etc.
    subhead = models.CharField(max_length=255, null=True, blank=True)

    tags = models.ManyToManyField(Tag, blank=True)

    _readonly = False  # Is this a read only model? (i.e. from elasticsearch)
    _cache = {}  # This is a cache for the content doctypes

    def __unicode__(self):
        return '%s: %s' % (self.__class__.__name__, self.title)

    def get_absolute_url(self):
        return reverse('content-detail-view', kwargs=dict(pk=self.pk, slug=self.slug))

    @property
    def byline(self):
        # If the subclass has customized the byline accessing, use that.
        if hasattr(self, 'get_byline'):
            return self.get_byline()

        # If we have an override byline, we'll use that first.
        if self._byline:
            return self._byline

        # If we have authors, just put them in a list
        if self.authors.exists():
            return ', '.join([user.get_full_name() for user in self.authors.all()])

        # Well, shit. I guess there's no byline.
        return None

    def extract_document(self):
        doc = super(Content, self).extract_document()
        doc.update({
            'published': self.published,
            'title': self.title,
            'slug': self.slug,
            'description': self.description,
            'image': self.image.name,
            'byline': self.byline,
            'subhead': self.subhead,
            'feature_type': self.feature_type,
            'feature_type.slug': slugify(self.feature_type)
        })
        if self.tags:
            doc['tags'] = [tag.extract_document() for tag in self.tags.all()]
        return doc

    @property
    def feature_type(self):
        # If the subclass has customized the feature_type accessing, use that.
        if hasattr(self, 'get_feature_type'):
            return self.get_feature_type()

        if self._feature_type:
            return self._feature_type

        return None

    @feature_type.setter
    def feature_type(self, value):
        if self._readonly:
            raise AttributeError('This content object is read only.')
        self._feature_type = value

    # class methods ##############################

    @classmethod
    def from_source(cls, _source):
        obj = cls(
            id=_source['id'],
            published=_source['published'],
            title=_source['title'],
            slug=_source['slug'],
            description=_source['description'],
            subhead=_source['subhead'],
            _feature_type=_source['feature_type'],
        )
        tags = [deserialize_polymorphic_model(tag_source) for tag_source in _source.get('tags', [])]
        obj.tags = tags
        return obj

    @classmethod
    def get_doctypes(cls):
        if len(cls._cache) == 0:
            for app in models.get_apps():
                for model in models.get_models(app, include_auto_created=True):
                    if isinstance(model(), Content):
                        cls._cache[model.get_mapping_type_name()] = model
        return cls._cache

    @classmethod
    def get_mapping_properties(cls):
        properties = super(Content, cls).get_mapping_properties()
        properties.update({
            'published': {'type': 'date'},
            'title': {'type': 'string'},
            'slug': {'type': 'string'},
            'description': {'type': 'string'},
            'image': {'type': 'integer'},
            'byline': {'type': 'string'},
            'feature_type': {
                'type': 'multi_field',
                'fields': {
                    'feature_type': {'type': 'string', 'index': 'not_analyzed'},
                    'slug': {'type': 'string', 'index': 'not_analyzed'}
                }
            },
            'tags': {
                'properties': Tag.get_mapping_properties()
            }
        })
        return properties

    @classmethod
    def get_serializer_class(cls):
        from .serializers import ContentSerializerReadOnly
        return ContentSerializerReadOnly

    @classmethod
    def get_writable_serializer_class(cls):
        from .serializers import ContentSerializer
        return ContentSerializer

    @classmethod
    def search(cls, **kwargs):
        """
        If ElasticSearch is being used, we'll use that for the query, and otherwise
        fall back to Django's .filter().

        Allowed params:

         * query
         * tag(s)
         * type(s)
         * feature_type(s)
         * published
        """
        index = settings.ES_INDEXES.get('default')
        results = ContentS().es(urls=settings.ES_URLS).indexes(index)
        if kwargs.get('pk'):
            try:
                pk = int(kwargs['pk'])
            except ValueError:
                pass
            else:
                results = results.query(id=kwargs['pk'])

        if kwargs.get('query'):
            results = results.query(_all__text_phrase=kwargs.get('query'))

        if kwargs.get('published', True):
            now = timezone.now()
            results = results.query(published__lte=now, must=True)

        for tag in kwargs.get('tags', []):
            tag_query_string = 'tags.slug:%s' % tag
            results = results.query(__query_string=tag_query_string)

        for feature_type in kwargs.get('feature_types', []):
            feature_type_query_string = 'feature_type.slug:%s' % feature_type
            results = results.query(__query_string=feature_type_query_string)

        types = kwargs.pop('types', [])
        if types:
            # only use valid subtypes
            results = results.doctypes(*[
                type_classname for type_classname in types \
                if type_classname in cls.get_doctypes()
            ])
        else:
            results = results.doctypes(*cls.get_doctypes().keys())

        return results.order_by('-published')

def content_tags_changed(sender, instance=None, action='', **kwargs):
    """Reindex content tags when they change."""
    es = get_es()
    indexes = settings.ES_INDEXES
    index = indexes['default']
    doc = {}
    doc['tags'] = [tag.extract_document() for tag in instance.tags.all()]
    es.update(index, instance.get_mapping_type_name(), instance.id, doc=doc, refresh=True)

models.signals.m2m_changed.connect(
    content_tags_changed,
    sender=Content.tags.through,
    dispatch_uid='content_tags_changed_signal'
)

<commit_msg>Made description and slug not required. Added `build_slug` method to `Content` which is run on save<commit_after>"""Base models for "Content", including the indexing and search features
that we want any piece of content to have."""

from django.conf import settings
from django.contrib.contenttypes.models import ContentType
from django.core.urlresolvers import reverse
from django.db import models
from django.db.backends import util
from django.template.defaultfilters import slugify
from django.utils import timezone

from bulbs.images.fields import RemoteImageField

from elasticutils import SearchResults, S
from elasticutils.contrib.django import get_es
from polymorphic import PolymorphicModel

class ReadonlyRelatedManager(object):
    """Replaces Django's RelatedMangers in read-only scenarios."""
    def __init__(self):
        self.data = []

    def __set__(self, obj, data):
        if data:
            if not isinstance(data, list):
                data = [data]
        else:
            data = []
        self.data = data

    def add(self, *args):
        self._on_bad_access()

    def all(self):
        return self.data

    def clear(self):
        self._on_bad_access()

    def remove(self, *args):
        self._on_bad_access()

    def _on_bad_access(self):
        raise TypeError('%s is read-only' % self.__class__.__name__)

def readonly_content_factory(model):
    """This is used to generate a class for a piece of content that we
    have retrieved from ElasticSearch. This content will have a number
    of "read-only" fields."""

    class Meta:
        proxy = True
        app_label = model._meta.app_label

    name = '%s_Readonly' % model.__name__
    name = util.truncate_name(name, 80, 32)

    overrides = {
        'Meta': Meta,
        '__module__': model.__module__,
        '_readonly': True,
        'tags': ReadonlyRelatedManager()
    }
    # TODO: Add additional deferred fields here, just as placeholders.
    return type(str(name), (model,), overrides)

def deserialize_polymorphic_model(data):
    """Deserializes simple polymorphic models."""
    content_type = ContentType.objects.get_for_id(data['polymorphic_ctype_id'])
    if content_type:
        klass = content_type.model_class()
        instance = klass.from_source(data)
        return instance

class ContentSearchResults(SearchResults):
    def set_objects(self, results):
        self.objects = []
        readonly_cls = readonly_content_factory(Content)
        for result in results:
            obj = readonly_cls.from_source(result['_source'])
            self.objects.append(obj)

    def __iter__(self):
        return self.objects.__iter__()

class ContentS(S):

    def all(self):
        """
        Fixes the default `S.all` method given by elasticutils.
        `S` generally looks like django queryset but differs in
        a few ways, one of which is `all`. Django `QuerySet` just
        returns a clone for `all` but `S` wants to return all
        the documents. This makes `all` at least respect slices
        but the real fix is to probably make `S` work more like
        `QuerySet`.
        """
        if self.start == 0 and self.stop is None:
            # no slicing has occurred. let's get all of the records.
            count = self.count()
            return self[:count].execute()
        return self.execute()

    def get_results_class(self):
        """Returns the results class to use

        The results class should be a subclass of SearchResults.

        """
        if self.as_list or self.as_dict:
            return super(ContentS, self).get_results_class()

        return ContentSearchResults

class TagSearchResults(SearchResults):
    def set_objects(self, results):
        self.objects = [
            deserialize_polymorphic_model(result['_source']) for result in results
        ]

    def __iter__(self):
        return self.objects.__iter__()

class TagS(S):

    def get_results_class(self):
        """Returns the results class to use

        The results class should be a subclass of SearchResults.

        """
        if self.as_list or self.as_dict:
            return super(TagS, self).get_results_class()

        return TagSearchResults

class PolymorphicIndexable(object):
    """Base mixin for polymorphic indexin'"""
    def extract_document(self):
        return {
            'id': self.id,
            'polymorphic_ctype_id': self.polymorphic_ctype_id
        }

    def index(self, refresh=False):
        es = get_es(urls=settings.ES_URLS)
        index = settings.ES_INDEXES.get('default')
        es.index(
            index,
            self.get_mapping_type_name(),
            self.extract_document(),
            self.id,
            refresh=refresh
        )

    def save(self, index=True, refresh=False, *args, **kwargs):
        result = super(PolymorphicIndexable, self).save(*args, **kwargs)
        if index:
            self.index(refresh=refresh)
        return result

    @classmethod
    def get_mapping(cls):
        return {
            cls.get_mapping_type_name(): {
                'properties': cls.get_mapping_properties()
            }
        }

    @classmethod
    def get_mapping_properties(cls):
        return {
            'id': {'type': 'integer'},
            'polymorphic_ctype_id': {'type': 'integer'}
        }

    @classmethod
    def get_mapping_type_name(cls):
        return '%s_%s' % (cls._meta.app_label, cls.__name__.lower())

class Tag(PolymorphicIndexable, PolymorphicModel):
    """Model for tagging up Content."""
    name = models.CharField(max_length=255)
    slug = models.SlugField()

    _doctype_cache = {}

    def __unicode__(self):
        return '%s: %s' % (self.__class__.__name__, self.name)

    def save(self, *args, **kwargs):
        self.slug = slugify(self.name)
        return super(Tag, self).save(*args, **kwargs)

    def extract_document(self):
        doc = super(Tag, self).extract_document()
        doc.update({
            'name': self.name,
            'slug': self.slug
        })
        return doc

    @classmethod
    def from_source(cls, _source):
        return cls(
            id=_source['id'],
            polymorphic_ctype_id=_source['polymorphic_ctype_id'],
            name=_source['name'],
            slug=_source['slug']
        )

    @classmethod
    def get_doctypes(cls):
        if len(cls._doctype_cache) == 0:
            for app in models.get_apps():
                for model in models.get_models(app, include_auto_created=True):
                    if isinstance(model(), Tag):
                        cls._doctype_cache[model.get_mapping_type_name()] = model
        return cls._doctype_cache

    @classmethod
    def get_mapping_properties(cls):
        props = super(Tag, cls).get_mapping_properties()
        props.update({
            'name': {'type': 'string', 'index': 'not_analyzed'},
            'slug': {'type': 'string', 'index': 'not_analyzed'},
        })
        return props

    @classmethod
    def search(cls, **kwargs):
        """Search tags...profit."""
        index = settings.ES_INDEXES.get('default')
        results = TagS().es(urls=settings.ES_URLS).indexes(index)
        name = kwargs.pop('name', '')
        if name:
            results = results.query(name__prefix=name, boost=4, should=True).query(name__fuzzy={
                'value': name,
                'prefix_length': 1,
                'min_similarity': 0.35
            }, should=True)

        types = kwargs.pop('types', [])
        if types:
            # only use valid subtypes
            results = results.doctypes(*[
                type_classname for type_classname in kwargs['types'] \
                if type_classname in cls.get_doctypes()
            ])
        else:
            results = results.doctypes(*cls.get_doctypes().keys())
        return results

class Section(Tag):
    """Tag subclass which represents major sections of the site."""
    class Meta(Tag.Meta):
        proxy = True

class Content(PolymorphicIndexable, PolymorphicModel):
    """The base content model from which all other content derives."""
    published = models.DateTimeField(blank=True, null=True)
    title = models.CharField(max_length=512)
    slug = models.SlugField(blank=True, default='')
    description = models.TextField(max_length=1024, blank=True, default='')
    image = RemoteImageField(null=True, blank=True)

    authors = models.ManyToManyField(settings.AUTH_USER_MODEL)
    _byline = models.CharField(max_length=255, null=True, blank=True)  # This is an overridable field that is by default the author names
    _tags = models.TextField(null=True, blank=True)  # A return-separated list of tag names, exposed as a list of strings
    _feature_type = models.CharField(max_length=255, null=True, blank=True)  # "New in Brief", "Newswire", etc.
    subhead = models.CharField(max_length=255, null=True, blank=True)

    tags = models.ManyToManyField(Tag, blank=True)

    _readonly = False  # Is this a read only model? (i.e. from elasticsearch)
    _cache = {}  # This is a cache for the content doctypes

    def __unicode__(self):
        return '%s: %s' % (self.__class__.__name__, self.title)

    def get_absolute_url(self):
        return reverse('content-detail-view', kwargs=dict(pk=self.pk, slug=self.slug))

    @property
    def byline(self):
        # If the subclass has customized the byline accessing, use that.
        if hasattr(self, 'get_byline'):
            return self.get_byline()

        # If we have an override byline, we'll use that first.
        if self._byline:
            return self._byline

        # If we have authors, just put them in a list
        if self.authors.exists():
            return ', '.join([user.get_full_name() for user in self.authors.all()])

        # Well, shit. I guess there's no byline.
        return None

    def build_slug(self):
        return self.title

    def extract_document(self):
        doc = super(Content, self).extract_document()
        doc.update({
            'published': self.published,
            'title': self.title,
            'slug': self.slug,
            'description': self.description,
            'image': self.image.name,
            'byline': self.byline,
            'subhead': self.subhead,
            'feature_type': self.feature_type,
            'feature_type.slug': slugify(self.feature_type)
        })
        if self.tags:
            doc['tags'] = [tag.extract_document() for tag in self.tags.all()]
        return doc

    @property
    def feature_type(self):
        # If the subclass has customized the feature_type accessing, use that.
        if hasattr(self, 'get_feature_type'):
            return self.get_feature_type()

        if self._feature_type:
            return self._feature_type

        return None

    @feature_type.setter
    def feature_type(self, value):
        if self._readonly:
            raise AttributeError('This content object is read only.')
        self._feature_type = value

    def save(self, *args, **kwargs):
        self.slug = slugify(self.build_slug())

        return super(Content, self).save(*args, **kwargs)
    # class methods ##############################

    @classmethod
    def from_source(cls, _source):
        obj = cls(
            id=_source['id'],
            published=_source['published'],
            title=_source['title'],
            slug=_source['slug'],
            description=_source['description'],
            subhead=_source['subhead'],
            _feature_type=_source['feature_type'],
        )
        tags = [deserialize_polymorphic_model(tag_source) for tag_source in _source.get('tags', [])]
        obj.tags = tags
        return obj

    @classmethod
    def get_doctypes(cls):
        if len(cls._cache) == 0:
            for app in models.get_apps():
                for model in models.get_models(app, include_auto_created=True):
                    if isinstance(model(), Content):
                        cls._cache[model.get_mapping_type_name()] = model
        return cls._cache

    @classmethod
    def get_mapping_properties(cls):
        properties = super(Content, cls).get_mapping_properties()
        properties.update({
            'published': {'type': 'date'},
            'title': {'type': 'string'},
            'slug': {'type': 'string'},
            'description': {'type': 'string'},
            'image': {'type': 'integer'},
            'byline': {'type': 'string'},
            'feature_type': {
                'type': 'multi_field',
                'fields': {
                    'feature_type': {'type': 'string', 'index': 'not_analyzed'},
                    'slug': {'type': 'string', 'index': 'not_analyzed'}
                }
            },
            'tags': {
                'properties': Tag.get_mapping_properties()
            }
        })
        return properties

    @classmethod
    def get_serializer_class(cls):
        from .serializers import ContentSerializerReadOnly
        return ContentSerializerReadOnly

    @classmethod
    def get_writable_serializer_class(cls):
        from .serializers import ContentSerializer
        return ContentSerializer

    @classmethod
    def search(cls, **kwargs):
        """
        If ElasticSearch is being used, we'll use that for the query, and otherwise
        fall back to Django's .filter().

        Allowed params:

         * query
         * tag(s)
         * type(s)
         * feature_type(s)
         * published
        """
        index = settings.ES_INDEXES.get('default')
        results = ContentS().es(urls=settings.ES_URLS).indexes(index)
        if kwargs.get('pk'):
            try:
                pk = int(kwargs['pk'])
            except ValueError:
                pass
            else:
                results = results.query(id=kwargs['pk'])

        if kwargs.get('query'):
            results = results.query(_all__text_phrase=kwargs.get('query'))

        if kwargs.get('published', True):
            now = timezone.now()
            results = results.query(published__lte=now, must=True)

        for tag in kwargs.get('tags', []):
            tag_query_string = 'tags.slug:%s' % tag
            results = results.query(__query_string=tag_query_string)

        for feature_type in kwargs.get('feature_types', []):
            feature_type_query_string = 'feature_type.slug:%s' % feature_type
            results = results.query(__query_string=feature_type_query_string)

        types = kwargs.pop('types', [])
        if types:
            # only use valid subtypes
            results = results.doctypes(*[
                type_classname for type_classname in types \
                if type_classname in cls.get_doctypes()
            ])
        else:
            results = results.doctypes(*cls.get_doctypes().keys())

        return results.order_by('-published')

def content_tags_changed(sender, instance=None, action='', **kwargs):
    """Reindex content tags when they change."""
    es = get_es()
    indexes = settings.ES_INDEXES
    index = indexes['default']
    doc = {}
    doc['tags'] = [tag.extract_document() for tag in instance.tags.all()]
    es.update(index, instance.get_mapping_type_name(), instance.id, doc=doc, refresh=True)

models.signals.m2m_changed.connect(
    content_tags_changed,
    sender=Content.tags.through,
    dispatch_uid='content_tags_changed_signal'
)

<|endoftext|>


<commit_before>import frappe

def execute():
	frappe.conn.sql("""update `tabWebsite Sitemap` set idx=null""")
	for doctype in ["Blog Category", "Blog Post", "Web Page", "Website Group"]:
		frappe.conn.sql("""update `tab{}` set idx=null""".format(doctype))

	from frappe.website.doctype.website_sitemap_config.website_sitemap_config import rebuild_website_sitemap_config
	rebuild_website_sitemap_config()<commit_msg>fixed patch<commit_after>import frappe

def execute():
	frappe.reload_doc("website", "doctype", "blog_post")
	frappe.conn.sql("""update `tabWebsite Sitemap` set idx=null""")
	for doctype in ["Blog Category", "Blog Post", "Web Page", "Website Group"]:
		frappe.conn.sql("""update `tab{}` set idx=null""".format(doctype))

	from frappe.website.doctype.website_sitemap_config.website_sitemap_config import rebuild_website_sitemap_config
	rebuild_website_sitemap_config()<|endoftext|>