# 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
#

# pylint: disable=protected-access

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


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_RECORD': [None, {
        'NextService': [0x0, ['_SERVICE_HEADER']],
        'ServiceName': [0x8, ['pointer', ['UnicodeString', dict(length=512)]]],
        'DisplayName': [0xc, ['pointer', ['UnicodeString', dict(length=512)]]],
        'Order': [0x10, ['unsigned int']],
        'Tag': [0x18, ['array', 4, ['unsigned char']]],
        'DriverName': [0x24, ['pointer', ['UnicodeString', dict(
            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', dict(
            count=4,
            target='unsigned char'
            )]],
        'ServiceRecord': [0x10, ['Pointer', dict(
            target='_SERVICE_RECORD'
            )]],
        }],

    '_SERVICE_RECORD': [None, {
        'NextService': [0x0, ['Pointer', dict(
            target="_SERVICE_RECORD"
            )]],
        'ServiceName': [0x8, ['pointer', ['UnicodeString', dict(
            encoding='utf16', length=512)]]],

        'DisplayName': [0x10, ['Pointer', dict(
            target='UnicodeString',
            target_args=dict(length=512)
            )]],
        'Order': [0x18, ['unsigned int']],
        'Tag' : [0x20, ['Array', dict(
            count=4,
            target='unsigned char'
            )]],
        'DriverName': [0x30, ['Pointer', dict(
            target='UnicodeString',
            target_args=dict(
                length=256
                )
            )]],
        'ServiceProcess': [0x30, ['Pointer', dict(
            target='_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', dict(
            target='UnicodeString',
            target_args=dict(
                length=256
                )
            )]],
        'ProcessId': [0x18, ['unsigned int']],
        }],
}


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)


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


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

    def is_valid(self):
        "Check some fields for validity"
        return (super(_SERVICE_HEADER, self).is_valid() and
                self.ServiceRecord.is_valid())


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


_SERVICE_RECORD_VISTA_X64 = {
    '_SERVICE_RECORD': [None, {
        'NextService': [0x00, ['Pointer', dict(
            target='_SERVICE_RECORD'
            )]],

        'ServiceName': [0x08, ['pointer', ['UnicodeString', dict(
            length=512
            )]]],
        'DisplayName': [0x10, ['pointer', ['UnicodeString', dict(
            length=512
            )]]],
        'Order': [0x18, ['unsigned int']],
        'ServiceProcess': [0x28, ['pointer', ['_SERVICE_PROCESS']]],
        'DriverName': [0x28, ['Pointer', dict(
            target='UnicodeString',
            target_args=dict(
                length=256,
                )
            )]],

        'Type' : [0x30, ['Flags', {'bitmap': SERVICE_TYPE_FLAGS}]],
        'State': [0x34, ['Enumeration', dict(
            target='unsigned int',
            choices=SERVICE_STATE_ENUM
            )]],
        }],
    }


_SERVICE_RECORD_WIN81_X64 = {
    "_SERVICE_RECORD": [None, {
            "Tag": [0, ["String", dict(length=4)]], # Signature sErv
            'NextService': [0x8, ['Pointer', dict(
                target='_SERVICE_RECORD'
                )]],

            'ServiceName': [0x10, ['pointer', ['UnicodeString', dict(
                length=512
                )]]],
            'DisplayName': [0x18, ['pointer', ['UnicodeString', dict(
                length=512
                )]]],
            'Order': [0x20, ['unsigned int']],
            'ServiceProcess': [0x38, ['pointer', ['_SERVICE_PROCESS']]],
            'DriverName': [0x38, ['Pointer', dict(
                target='UnicodeString',
                target_args=dict(
                    length=256,
                    )
                )]],

            'Type' : [0x40, ['Flags', {'bitmap': SERVICE_TYPE_FLAGS}]],
            'State': [0x44, ['Enumeration', dict(
                target='unsigned int',
                choices=SERVICE_STATE_ENUM
                )]],
            }],

    '_SERVICE_PROCESS': [None, {
        'Tag': [0, ["String", dict(length=4)]], # Sc16
        'BinaryPath': [0x18, ['Pointer', dict(
            target='UnicodeString',
            target_args=dict(
                length=256
                )
            )]],
        'ProcessId': [0x20, ['unsigned int']],
        }],
    }


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

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

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

        # Vista 2008 and windows 7
        elif '6.0' <= version <= '6.2':
            profile.add_classes({
                    '_SERVICE_RECORD': _SERVICE_RECORD_RECENT,
                    '_SERVICE_HEADER': _SERVICE_HEADER,
                    })
            profile.add_constants(ServiceTag="serH")

            if profile.metadata("arch") == "I386":
                profile.add_overlay(_SERVICE_RECORD_VISTA_X86)
            else:
                profile.add_overlay(_SERVICE_RECORD_VISTA_X64)

        # Windows 8.1
        elif '6.3' <= version <= '6.3':
            profile.add_classes({
                '_SERVICE_RECORD': _SERVICE_RECORD_RECENT,
                '_SERVICE_HEADER': _SERVICE_HEADER,
                })
            profile.add_constants(ServiceTag="serH")

            if profile.metadata("arch") == "I386":
                profile.add_overlay(_SERVICE_RECORD_VISTA_X86)
            else:
                profile.add_overlay(_SERVICE_RECORD_WIN81_X64)

        else:
            raise RuntimeError(
                "Unsupported windows version. Please file a bug.")



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

    def __init__(self, **kwargs):
        super(SvcRecordScanner, self).__init__(**kwargs)
        self.checks = [
            ('StringCheck', dict(
                    needle=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 = [
            ('StringCheck', dict(
                    needle=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.walk_list("NextService"):
                    if record.is_valid() and record not in self.records:
                        self.records.add(record)

                        yield record


class SvcScan(plugin.KernelASMixin, 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,
                    session=self.session)
            else:
                # Windows Vista, 2008, and 7 use the _SERVICE_HEADER
                scanner = SvcHeaderScanner(
                    task=task, process_profile=self.profile,
                    session=self.session)

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

    def render(self, renderer):
        renderer.table_header([
            ("Offset", "offset", "[addrpad]"),
            ("Order", "order", "5"),
            ("PID", "pid", "4"),
            ("Service Name", "service", "[wrap:30]"),
            ("Display Name", "display_name", "[wrap:40]"),
            ("Service Type", "type", "[wrap:30]"),
            ("Service State", "state", "15"),
            ("Binary Path", "binary_path", "")])

        for rec in self.calculate():
            renderer.table_row(
                rec,
                rec.Order,
                rec.Pid,
                rec.ServiceName.deref(),
                rec.DisplayName.deref(),
                rec.Type,
                rec.State,
                rec.Binary)
