# Rekall Memory Forensics
#
# Copyright 2014 Google Inc. All Rights Reserved.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or (at
# your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
#

"""
The Rekall Entity Layer.
"""
__author__ = "Adam Sindelar <adamsh@google.com>"

import re

from rekall import registry


class EntityCollector(object):
    """Base class for entity collectors.

    EntityCollector subclasses need to override the 'collect' method and the
    'collects' ivar. EntityCollector takes care of registering collected
    entities with the EntityManager, deciding whether the collector should be
    called for manager-driven queries and ensuring collection is only done once
    per cache lifetime.
    """

    collects = []  # Subclasses must override. See 'can_collect' below.
    capabilities = None  # Generated by __ensure_compile_capabilities.

    __metaclass__ = registry.MetaclassRegistry
    __abstract = True

    DELIMITER = re.compile(r"[\/=]")

    def __init__(self, entity_manager=None):
        self.entity_manager = entity_manager
        self.session = entity_manager.session

    @property
    def profile(self):
        return self.session.profile

    @property
    def name(self):
        return getattr(self, "_name", type(self).__name__)

    @classmethod
    def __ensure_compile_capabilities(cls):
        """Compiles self.collects into self.capabilities."""
        if cls.capabilities:
            return

        cls.capabilities = [cls.DELIMITER.split(x) for x in cls.collects]

    @property
    def is_collected(self):
        """Set to True after self.collect runs with no hints.

        Running collect with a non-null hint will not flip this."""
        return self.name in self.entity_manager.finished_collectors

    def collect(self, hint=None):
        """Override to yield components - analogous to 'calculate', but typed.

        Subclasses should override this to yield components that represent the
        data they collect. For example, a pslist collector would yield Process
        components. Only components listed in self.collects should be yielded
        and no other.

        Output format:
        ==============

        In the simplest case, yield one component at a time. To indicate two
        (or more) components are related, yield both in a list (for example,
        Process and the MemoryObject wrapping the proc struct representing it
        in kernel).

        Unless an identity is specified explicitly, one will be created
        automatically by taking the first field of the first component and
        instantiating an identity with it.

        To specify an identity explicitly, yield the Entity component as first
        in the list and provide the identity in its 'identity' field.

        Examples:
        =========

        # Yielding just one component at a time:
        yield manager.Process(
            pid=proc.p_pid,
            ...)

        # Yielding the process component and the BSD proc struct:
        yield [
            manager.Process(pid=proc.p_pid, ...),
            manager.MemoryObject(base_object=proc, ...),
        ]

        # In both above examples, the identity will be "Process/pid=(PID here)."
        # Specifying an explicit identity would look like this:
        yield [
            manager.Entity(identity=ProcessIdentity(...)),
            manager.Process(pid=proc.p_pid, ...),
            manager.MemoryObject(base_object=proc, ...),
        ]

        Optimization hints:
        ===================

        A hint may be passed to the collect method, in which case the collector
        may elect to only collect some information it's capable of collecting
        and not all. For example, a hint to the pslist collector may specify
        that the caller only cares about processes listed by walking the
        process head list. The format of the hint is currently not specified.
        """
        pass

    @classmethod
    def can_collect(cls, wanted):
        """Should this collector be called to collect the wanted?

        The format for specifying what a collector can collect is as follows:
        COMPONENT_NAME[/ATTRIBUTE_NAME[=VALUE]]

        Alternatively, wanted may be given as tuple of:
        ("COMPONENT/ATTRIBUTE", VALUE)

        A collector may specify as many such capabilities as needed in
        the self.collects ivar. For example, the pslist collector might specify:

        collects = ["Process", "MemoryObject/type=proc"]

        The wanted parameter is given using the same format and will be tested
        against each capability until one matches, with the test being only
        as specific as the query. For example, using the above definition, True
        would be returned for any of the following values of wanted:

        ("Process", "Process/state=running", "MemoryObject",
        "MemoryObject/type=proc")

        ...but none of the following:

        ("AllocationZone", "MemoryObject/type=zone")
        """
        if isinstance(wanted, tuple):
            # It's legal to supply wanted as (attribute, value).
            attribute, value = wanted
            wanted_pieces = attribute.split("/")
            wanted_pieces.append(value)
        else:
            # Otherwise it's a string.
            wanted_pieces = cls.DELIMITER.split(wanted)

        if wanted_pieces[0] in ("*", "Entity"):
            return True

        cls.__ensure_compile_capabilities()
        count_wanted = len(wanted_pieces)

        for capability in cls.capabilities:
            limit = min(len(capability), count_wanted)
            if wanted_pieces[:limit] == capability[:limit]:
                return True

        return False

    @classmethod
    def is_active(cls, session):
        _ = session
        return False

    @property
    def entities(self):
        """Returns this collector's entities."""
        return self.session.entities.find_by_collector(self.name)
