# Rekall Memory Forensics
# Copyright (c) 2010, 2011, 2012 Michael Ligh <michael.ligh@mnin.org>
# Copyright 2013 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
#

from rekall import obj
from rekall.plugins.windows import common
from rekall.plugins.windows import vadinfo


#--------------------------------------------------------------------------------
# vtypes
#--------------------------------------------------------------------------------

SERVICE_TYPE_FLAGS = {
            'SERVICE_KERNEL_DRIVER': 0,
            'SERVICE_FILE_SYSTEM_DRIVER': 1,
            'SERVICE_WIN32_OWN_PROCESS': 4,
            'SERVICE_WIN32_SHARE_PROCESS': 5,
            'SERVICE_INTERACTIVE_PROCESS': 8}

SERVICE_STATE_ENUM = {
            1: 'SERVICE_STOPPED',
            2: 'SERVICE_START_PENDING',
            3: 'SERVICE_STOP_PENDING',
            4: 'SERVICE_RUNNING',
            5: 'SERVICE_CONTINUE_PENDING',
            6: 'SERVICE_PAUSE_PENDING',
            7: 'SERVICE_PAUSED'}

svcscan_base_x86 = {
    '_SERVICE_HEADER': [ None, {
    'Tag': [ 0x0, ['array', 4, ['unsigned char']]],
    'ServiceRecord': [ 0xC, ['pointer', ['_SERVICE_RECORD']]],
    } ],
    '_SERVICE_LIST_ENTRY' : [ 0x8, {
    'Blink' : [ 0x0, ['pointer', ['_SERVICE_RECORD']]],
    'Flink' : [ 0x4, ['pointer', ['_SERVICE_RECORD']]],
    } ],
    '_SERVICE_RECORD' : [ None, {
    'ServiceList' : [ 0x0, ['_SERVICE_LIST_ENTRY']],
    'ServiceName' : [ 0x8, ['pointer', ['UnicodeString', dict(
                            encoding = 'utf16', length = 512)]]],
    'DisplayName' : [ 0xc, ['pointer', ['UnicodeString', dict(
                            encoding = 'utf16', length = 512)]]],
    'Order' : [ 0x10, ['unsigned int']],
    'Tag' : [ 0x18, ['array', 4, ['unsigned char']]],
    'DriverName' : [ 0x24, ['pointer', ['UnicodeString', dict(
                            encoding = 'utf16', length = 256)]]],
    'ServiceProcess' : [ 0x24, ['pointer', ['_SERVICE_PROCESS']]],
    'Type' : [ 0x28, ['Flags', {'bitmap': SERVICE_TYPE_FLAGS}]],
    'State' : [ 0x2c, ['Enumeration', dict(target='long',
                                           choices=SERVICE_STATE_ENUM)]],
    }],
    '_SERVICE_PROCESS' : [ None, {
    'BinaryPath' : [ 0x8, ['pointer', ['UnicodeString', dict(
                            encoding='utf16', length=256)]]],
    'ProcessId' : [ 0xc, ['unsigned int']],
    } ],
}

svcscan_base_x64 = {
    '_SERVICE_HEADER': [ None, {
    'Tag': [ 0x0, ['array', 4, ['unsigned char']]],
    'ServiceRecord': [ 0x10, ['pointer', ['_SERVICE_RECORD']]],
    } ],
    '_SERVICE_LIST_ENTRY' : [ 0x8, {
    'Blink' : [ 0x0, ['pointer', ['_SERVICE_RECORD']]],
    'Flink' : [ 0x10, ['pointer', ['_SERVICE_RECORD']]],
    } ],
    '_SERVICE_RECORD' : [ None, {
    'ServiceList' : [ 0x0, ['_SERVICE_LIST_ENTRY']],
    'ServiceName' : [ 0x8, ['pointer', ['UnicodeString', dict(
                            encoding='utf16', length=512)]]],
    'DisplayName' : [ 0x10, ['pointer', ['UnicodeString', dict(
                            encoding='utf16', length=512)]]],
    'Order' : [ 0x18, ['unsigned int']],
    'Tag' : [ 0x20, ['array', 4, ['unsigned char']]],
    'DriverName' : [ 0x30, ['pointer', ['UnicodeString', dict(
                            encoding='utf16', length=256)]]],
    'ServiceProcess' : [ 0x30, ['pointer', ['_SERVICE_PROCESS']]],
    'Type' : [ 0x38, ['Flags', {'bitmap': SERVICE_TYPE_FLAGS}]],
    'State' : [ 0x3C, ['Enumeration', dict(
                        target='long', choices=SERVICE_STATE_ENUM)]],
    } ],
    '_SERVICE_PROCESS': [ None, {
    'BinaryPath': [ 0x10, ['pointer', ['UnicodeString', dict(
                            encoding='utf16', length=256)]]],
    'ProcessId': [ 0x18, ['unsigned int']],
    } ],
}

#--------------------------------------------------------------------------------
# object Classes
#--------------------------------------------------------------------------------

class _SERVICE_RECORD_LEGACY(obj.Struct):
    "Service records for XP/2003 x86 and x64"

    @property
    def Binary(self):
        "Return the binary path for a service"

        # No path in memory for services that aren't running
        # (if needed, query the registry key)
        if str(self.State) != 'SERVICE_RUNNING':
            return obj.NoneObject("No path, service isn't running")

        # Depending on whether the service is for a process
        # or kernel driver, the binary path is stored differently
        if 'PROCESS' in str(self.Type):
            return self.ServiceProcess.BinaryPath.dereference()
        else:
            return self.DriverName.dereference()

    @property
    def Pid(self):
        "Return the process ID for a service"

        if str(self.State) == 'SERVICE_RUNNING':
            if 'PROCESS' in str(self.Type):
                return self.ServiceProcess.ProcessId

        return obj.NoneObject("Cannot get process ID")

    def is_valid(self):
        "Check some fields for validity"
        return (super(_SERVICE_RECORD_LEGACY, self).is_valid() and
                self.Order > 0 and self.Order < 0xFFFF)

    def traverse(self):

        rec = self # Include this object in the list
        while rec and rec.is_valid():
            yield rec
            rec = rec.ServiceList.Blink.dereference()

class _SERVICE_RECORD_RECENT(_SERVICE_RECORD_LEGACY):
    "Service records for 2008, Vista, 7 x86 and x64"

    def traverse(self):
        """Generator that walks the singly-linked list"""

        yield self # Include this object in the list

        # Make sure we dereference these pointers, or the
        # is_valid() checks will apply to the pointer and
        # not the _SERVICE_RECORD object as intended.
        rec = self.PrevEntry.dereference()
        while rec and rec.is_valid():
            yield rec
            rec = rec.PrevEntry.dereference()

class _SERVICE_HEADER(obj.Struct):
    "Service headers for 2008, Vista, 7 x86 and x64"

    def is_valid(self):
        "Check some fields for validity"
        return (obj.Struct.is_valid(self) and
                    self.ServiceRecord.is_valid() and
                    self.ServiceRecord.Order < 0xFFFF)


_SERVICE_RECORD_VISTA_X86 = {
    '_SERVICE_RECORD': [None, {
            'PrevEntry': [ 0x0, ['pointer', ['_SERVICE_RECORD']]],
            'ServiceName': [ 0x4, ['pointer', ['UnicodeString', dict(
                            encoding = 'utf16', length = 512)]]],
            'DisplayName': [ 0x8, ['pointer', ['UnicodeString', dict(
                            encoding = 'utf16', length = 512)]]],
            'Order': [ 0xC, ['unsigned int']],
            'ServiceProcess': [ 0x1C, ['pointer', ['_SERVICE_PROCESS']]],
            'DriverName': [ 0x1C, ['pointer', ['UnicodeString', dict(
                            encoding = 'utf16', length = 256)]]],
            'Type' : [ 0x20, ['Flags', {'bitmap': SERVICE_TYPE_FLAGS}]],
            'State': [ 0x24, ['Enumeration', dict(
                        target = 'long', choices = SERVICE_STATE_ENUM)]],
            }],
    }


_SERVICE_RECORD_VISTA_X64 = {
    '_SERVICE_RECORD': [None, {
            'PrevEntry': [ 0x0, ['pointer', ['_SERVICE_RECORD']]],
            'ServiceName': [ 0x8, ['pointer', ['UnicodeString', dict(
                            encoding = 'utf16', length = 512)]]],
            'DisplayName': [ 0x10, ['pointer', ['UnicodeString', dict(
                            encoding = 'utf16', length = 512)]]],
            'Order': [ 0x18, ['unsigned int']],
            'ServiceProcess': [ 0x28, ['pointer', ['_SERVICE_PROCESS']]],
            'DriverName': [ 0x28, ['pointer', ['UnicodeString', dict(
                            encoding = 'utf16', length = 256)]]],
            'Type' : [ 0x30, ['Flags', {'bitmap': SERVICE_TYPE_FLAGS}]],
            'State': [ 0x34, ['Enumeration', dict(target = 'long', choices = SERVICE_STATE_ENUM)]],
            }],
    }


class ServiceModification(obj.ProfileModification):
    """A modification for the service control manager."""

    @classmethod
    def modify(cls, profile):
        if profile.metadata("memory_model") == "32bit":
            profile.add_overlay(svcscan_base_x86)
        else:
            # 32bit Vista profiles
            profile.add_overlay(svcscan_base_x64)

        # Windows XP, 2003
        if profile.metadata("major") < 6:
            profile.add_classes({
                    '_SERVICE_RECORD': _SERVICE_RECORD_LEGACY,
                    '_SERVICE_HEADER': _SERVICE_HEADER,
                    })
            profile.add_constants(ServiceTag="sErv")

        # Vista 2008 and windows 7
        elif profile.metadata("major") >= 6:
            profile.add_classes({
                    '_SERVICE_RECORD': _SERVICE_RECORD_RECENT,
                    '_SERVICE_HEADER': _SERVICE_HEADER,
                    })
            profile.add_constants(ServiceTag="serH")

            if profile.metadata("memory_model") == "32bit":
                profile.add_overlay(_SERVICE_RECORD_VISTA_X86)
            else:
                profile.add_overlay(_SERVICE_RECORD_VISTA_X64)


#--------------------------------------------------------------------------------
# svcscan plugin
#--------------------------------------------------------------------------------

class SvcRecordScanner(vadinfo.VadScanner):
    """A scanner for the service tags."""

    def __init__(self, **kwargs):
        super(SvcRecordScanner, self).__init__(**kwargs)
        self.checks = [
            ('MultiStringFinderCheck', dict(
                    needles=[self.profile.get_constant("ServiceTag")])),
            ]
        self.tag_offset = self.profile.get_obj_offset('_SERVICE_RECORD', 'Tag')

    def scan(self, **kwargs):
        for hit in super(SvcRecordScanner, self).scan(**kwargs):
            svc_record = self.profile._SERVICE_RECORD(
                vm=self.address_space, offset=hit - self.tag_offset)

            if svc_record.is_valid():
                yield svc_record


class SvcHeaderScanner(vadinfo.VadScanner):
    """A scanner for the service tags."""

    def __init__(self, **kwargs):
        super(SvcHeaderScanner, self).__init__(**kwargs)
        self.checks = [
            ('MultiStringFinderCheck', dict(
                    needles=[self.profile.get_constant("ServiceTag")])),
            ]

        # On systems more recent than XP/2003, the serH marker doesn't
        # find *all* services, but the ones it does find have linked
        # lists to the others. We use this variable to track which
        # ones we've seen so as to not yield duplicates.
        self.records = set()

    def scan(self, **kwargs):
        for hit in super(SvcHeaderScanner, self).scan(**kwargs):
            svc_header = self.profile._SERVICE_HEADER(
                vm=self.address_space, offset=hit)

            if svc_header.is_valid():
                for record in svc_header.ServiceRecord.traverse():
                    # Since we walk the s-list backwards, if we've seen an object,
                    # then we've also seen all objects that exist before it, thus we
                    # can break at that time.
                    if record in self.records:
                        break

                    self.records.add(record)

                    yield record


class SvcScan(common.KDBGMixin, common.AbstractWindowsCommandPlugin):
    "Scan for Windows services"

    __name = "svcscan"

    def __init__(self, scan_in_kernel_address_space=False, **kwargs):
        """Scan for callbacks.

        Args:
           scan_in_kernel_address_space: If False we will use the physical
           address space for scanning, while if true we scan in the kernel
           address space.
        """
        super(SvcScan, self).__init__(**kwargs)
        self.scan_in_kernel_address_space = scan_in_kernel_address_space

        # Update the profile.
        self.profile = ServiceModification(self.profile)

    def calculate(self):
        # Get the version we're analyzing
        version = self.profile.metadatas('major', 'minor')

        pslist = self.session.plugins.pslist(proc_regex="services.exe")
        for task in pslist.filter_processes():
            # Process AS must be valid
            process_space = task.get_process_address_space()
            if process_space == None:
                continue

            # XP/2003 use the _SERVICE_RECORD object.
            if version <= (5,2):
                scanner = SvcRecordScanner(
                    task=task, process_profile=self.profile)
            else:
                # Windows Vista, 2008, and 7 use the _SERVICE_HEADER
                scanner = SvcHeaderScanner(
                    task=task, process_profile=self.profile)

            # Find all instances of the record tag
            for record in scanner.scan():
                yield record

    def render(self, renderer):
        for rec in self.calculate():
            # This can't possibly look neat in a table with columns...
            renderer.record(
                (("Offset", "offset", "{0:#x}", rec.obj_offset),
                 ("Order", "order", "{0}", rec.Order),
                 ("Process ID", "pid", "{0}", rec.Pid),
                 ("Service Name", "service", "{0}", rec.ServiceName.deref()),
                 ("Display Name", "display_name", "{0}", rec.DisplayName.deref()),
                 ("Service Type", "type", "{0}", rec.Type),
                 ("Service State", "state", "{0}", rec.State),
                 ("Binary Path", "binary_path", "{0}", rec.Binary)))

