# Rekall Memory Forensics
#
# Copyright 2013 Google Inc. All Rights Reserved.
#
# Authors:
# Michael Hale Ligh <michael.ligh@mnin.org>
#
# Contributors/References:
#   Richard Stevens and Eoghan Casey
#   Extracting Windows Cmd Line Details from Physical Memory.
#   http://ww.dfrws.org/2010/proceedings/stevens.pdf
#   Michael Cohen <scudette@gmail.com>
#
# 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 builtins import bytes

from rekall import obj
from rekall.plugins.overlays import basic
from rekall.plugins.windows import common
from rekall.plugins.windows import vadinfo
from rekall.plugins.overlays.windows import pe_vtypes
from rekall_lib import utils


MAX_HISTORY_DEFAULT = 50
HISTORY_BUFFERS_DEFAULTS = 4

# The following all use these standard types.
common_types = {
    '_LIST_ENTRY' : [0x8, {
        'Flink' : [0x0, ['pointer', ['_LIST_ENTRY']]],
        'Blink' : [0x4, ['pointer', ['_LIST_ENTRY']]],
        }]}

common_types_64 = {
    '_LIST_ENTRY' : [0x10, {
        'Flink' : [0x0, ['pointer', ['_LIST_ENTRY']]],
        'Blink' : [0x8, ['pointer', ['_LIST_ENTRY']]],
        }]}

# Windows 7 Types from conhost.exe
conhost_types_x86 = {
    '_COMMAND': [None, {
        'CmdLength': [0x00, ['unsigned short']],
        'Cmd' : [0x02, ['UnicodeString', dict(
            encoding='utf16',
            length=lambda x: x.CmdLength
            )]],
        }],
    '_COMMAND_HISTORY': [None, {
        'ListEntry': [0x00, ['_LIST_ENTRY']],
        'Flags' : [0x08, ['Flags', dict(
            bitmap={
                'Allocated': 0,
                'Reset': 1
                }
            )]],
        'Application': [0x0C, ['Pointer', dict(
            target='UnicodeString',
            target_args=dict(
                encoding='utf16',
                length=256
                )
            )]],
        'CommandCount': [0x10, ['short']],
        'LastAdded': [0x12, ['short']],
        'LastDisplayed': [0x14, ['short']],
        'FirstCommand': [0x16, ['short']],
        'CommandCountMax': [0x18, ['short']],
        'ProcessHandle': [0x1C, ['unsigned int']],
        'PopupList': [0x20, ['_LIST_ENTRY']],
        'CommandBucket': [0x28, ['Array', dict(
            count=lambda x: x.CommandCount,
            target='Pointer',
            target_args=dict(
                target='_COMMAND'
                )
            )]],
        }],

    '_ALIAS': [None, {
        'ListEntry': [0x00, ['_LIST_ENTRY']],
        'SourceLength': [0x08, ['unsigned short']],
        'TargetLength': [0x0A, ['unsigned short']],
        'Source': [0x0C, ['Pointer', dict(
            target='UnicodeString',
            target_args=dict(
                encoding='utf16',
                length=lambda x: x.SourceLength * 2
                )
            )]],

        'Target': [0x10, ['Pointer', dict(
            target='UnicodeString',
            target_args=dict(
                encoding='utf16',
                length=lambda x: x.TargetLength * 2
                )
            )]],
        }],
    '_EXE_ALIAS_LIST' : [None, {
        'ListEntry': [0x00, ['_LIST_ENTRY']],
        'ExeLength': [0x08, ['unsigned short']],
        'ExeName': [0x0C, ['Pointer', dict(
            target='UnicodeString',
            target_args=dict(
                encoding='utf16',
                length=lambda x: x.ExeLength * 2
                )
            )]],
        'AliasList': [0x10, ['_LIST_ENTRY']],
        }],

    '_POPUP_LIST' : [None, {
        'ListEntry' : [0x00, ['_LIST_ENTRY']],
        }],

    '_CONSOLE_INFORMATION': [None, {
        'CurrentScreenBuffer': [0x98, ['pointer', ['_SCREEN_INFORMATION']]],
        'ScreenBuffer': [0x9C, ['pointer', ['_SCREEN_INFORMATION']]],
        'HistoryList': [0xD4, ['_LIST_ENTRY']],
        'ProcessList': [0x18, ['_LIST_ENTRY']], # SrvGetConsoleProcessList()
        'ExeAliasList': [0xDC, ['_LIST_ENTRY']], # GetConsoleAliasExes()

        # GetConsoleHistoryInfo()
        'HistoryBufferCount': [0xE4, ['unsigned short']],

        # GetConsoleHistoryInfo()
        'HistoryBufferMax': [0xE6, ['unsigned short']],

        'CommandHistorySize': [0xE8, ['unsigned short']],
        'OriginalTitle': [0xEC, ['Pointer', dict(
            target='UnicodeString',
            target_args=dict(
                encoding='utf16',
                length=256
                )
            )]], # GetConsoleOriginalTitle()

        'Title': [0xF0, ['Pointer', dict(
            target='UnicodeString',
            target_args=dict(
                encoding='utf16',
                length=256
                )
            )]], # GetConsoleTitle()
        }],

    '_CONSOLE_PROCESS': [None, {
        'ListEntry': [0x00, ['_LIST_ENTRY']],
        'ProcessHandle': [0x8, ['unsigned int']],
        }],
    '_SCREEN_INFORMATION': [None, {
        'ScreenX': [0x08, ['short']],
        'ScreenY': [0x0A, ['short']],
        'Rows': [0x3C, ['Pointer', dict(
            target='Array',
            target_args=dict(
                count=lambda x: x.ScreenY,
                target='_ROW'
                )
            )]],
        'Next': [0xDC, ['Pointer', dict(target='_SCREEN_INFORMATION')]],
        }],
    '_ROW': [0x1C, {
        'Chars': [0x08, ['Pointer', dict(
            target='UnicodeString',
            target_args=dict(
                encoding='utf16',
                length=lambda x: x.obj_parent.ScreenX * 2,
                )
            )]],
        }],
    }

# Windows 7 Types from conhost.exe
conhost_types_x64 = {
    '_COMMAND': [None, {
        'CmdLength': [0x00, ['unsigned short']],
        'Cmd' : [0x02, ['UnicodeString', dict(
            encoding='utf16',
            length=lambda x: x.CmdLength)]],
        }],
    '_COMMAND_HISTORY': [None, {
        'ListEntry': [0x00, ['_LIST_ENTRY']],

        # AllocateCommandHistory()
        'Flags' : [0x10, ['Flags', {'bitmap': {'Allocated': 0, 'Reset': 1}}]],

        # AllocateCommandHistory()
        'Application': [0x18, ['Pointer', dict(
            target='UnicodeString',
            target_args=dict(
                encoding='utf16',
                length=256
                )
            )]],
        'CommandCount': [0x20, ['short']],
        'LastAdded': [0x22, ['short']],
        'LastDisplayed': [0x24, ['short']],
        'FirstCommand': [0x26, ['short']],
        'CommandCountMax': [0x28, ['short']], # AllocateCommandHistory()
        'ProcessHandle': [0x30, ['address']], # AllocateCommandHistory()
        'PopupList': [0x38, ['_LIST_ENTRY']], # AllocateCommandHistory()
        'CommandBucket': [0x48, ['Array', dict(
            count=lambda x: x.CommandCount,
            target='Pointer',
            target_args=dict(
                target='_COMMAND')
            )]],
        }],
    '_ALIAS': [None, {
        'ListEntry': [0x00, ['_LIST_ENTRY']],
        'SourceLength': [0x10, ['unsigned short']], # AddAlias()
        'TargetLength': [0x12, ['unsigned short']], # AddAlias()

        # reversed from AddAlias()
        'Source': [0x18, ['pointer', ['UnicodeString', dict(
            encoding='utf16',
            length=lambda x: x.SourceLength * 2)]]],

        'Target': [0x20, ['pointer', ['UnicodeString', dict(
            encoding='utf16',
            length=lambda x: x.TargetLength * 2
            )]]], # AddAlias()
        }],
    '_EXE_ALIAS_LIST' : [None, {
        'ListEntry': [0x00, ['_LIST_ENTRY']],
        'ExeLength': [0x10, ['unsigned short']], # AddExeAliasList()

        # AddExeAliasList()
        'ExeName': [0x18, ['pointer', ['UnicodeString', dict(
            encoding='utf16',
            length=lambda x: x.ExeLength * 2)]]],

        'AliasList': [0x20, ['_LIST_ENTRY']], # AddExeAliasList()
        }],
    '_POPUP_LIST' : [None, {
        'ListEntry' : [0x00, ['_LIST_ENTRY']],
        }],
    '_CONSOLE_INFORMATION': [None, {
        # MOV RCX, [RIP+0x25bb5]   conhost!gConsoleInformation+0x28
        'ProcessList': [0x28, ['_LIST_ENTRY']],
        'CurrentScreenBuffer': [0xE0, ['Pointer', dict(
            target='_SCREEN_INFORMATION'
            )]], # AllocateConsole()

        'ScreenBuffer': [0xE8, ['Pointer', dict(
            target='_SCREEN_INFORMATION'
            )]], # AllocateConsole()

        'HistoryList': [0x148, ['_LIST_ENTRY']], # AllocateCommandHistory()
        'ExeAliasList': [0x148, ['_LIST_ENTRY']], # SrvGetConsoleAliasExes()
        'HistoryBufferCount': [0x168, ['unsigned short']], # AllocateConsole()
        'HistoryBufferMax': [0x16A, ['unsigned short']], # AllocateConsole()
        'CommandHistorySize': [0x16C, ['unsigned short']], # AllocateConsole()

        'OriginalTitle': [0x170, ['Pointer', dict(
            target='UnicodeString',
            target_args=dict(
                encoding='utf16',
                length=256
                )
            )]], # SrvGetConsoleTitle()

        'Title': [0x178, ['Pointer', dict(
            target='UnicodeString',
            target_args=dict(
                encoding='utf16',
                length=256
                )
            )]], # SrvGetConsoleTitle()
        }],
    '_CONSOLE_PROCESS': [None, {
        'ListEntry': [0x00, ['_LIST_ENTRY']],
        'ProcessHandle': [0x10, ['unsigned int']], # FindProcessInList()
        }],

    '_SCREEN_INFORMATION': [None, {
        'ScreenX': [8, ['short']],
        'ScreenY': [10, ['short']],
        'Rows': [0x48, ['Pointer', dict(
            target="Array",
            target_args=dict(
                count=lambda x: x.ScreenY,
                target='_ROW'
                )
            )]],

        'Next': [0x128, ['pointer', ['_SCREEN_INFORMATION']]],
        }],
    '_ROW': [0x28, {
        'Chars': [0x08, ['pointer', ['UnicodeString', dict(
            encoding='utf16',
            length=lambda x: x.obj_parent.obj_parent.ScreenX * 2,
        )]]],
        }],
    }

# Windows XP, 2003, 2008, Vista from winsrv.dll
winsrv_types_x86 = {
    '_COMMAND': [None, {
        'CmdLength': [0x00, ['unsigned short']],
        'Cmd' : [0x02, ['UnicodeString', dict(
            encoding='utf16',
            length=lambda x: x.CmdLength
            )]],
        }],
    '_COMMAND_HISTORY': [None, {
        'Flags' : [0x00, ['Flags', {
            'bitmap': {
                'Allocated': 0,
                'Reset': 1
                }
            }]],
        'ListEntry': [0x04, ['_LIST_ENTRY']],
        'Application': [0x0C, ['pointer', ['UnicodeString', dict(
            encoding='utf16', length=256)]]],
        'CommandCount': [0x10, ['short']],
        'LastAdded': [0x12, ['short']],
        'LastDisplayed': [0x14, ['short']],
        'FirstCommand': [0x16, ['short']],
        'CommandCountMax': [0x18, ['short']],
        'ProcessHandle': [0x1C, ['unsigned int']],
        'PopupList': [0x20, ['_LIST_ENTRY']],
        'CommandBucket': [0x28, ['Array', dict(
            count=lambda x: x.CommandCount,
            target='Pointer',
            target_args=dict(
                target='_COMMAND'
                )
            )]],
        }],

    '_ALIAS': [None, {
        'ListEntry': [0x00, ['_LIST_ENTRY']],
        'SourceLength': [0x08, ['unsigned short']],
        'TargetLength': [0x0A, ['unsigned short']],
        'Source': [0x0C, ['pointer', ['UnicodeString', dict(
            encoding='utf16',
            length=lambda x: x.SourceLength * 2
            )]]],
        'Target': [0x10, ['pointer', ['UnicodeString', dict(
            encoding='utf16',
            length=lambda x: x.TargetLength * 2)]]],
        }],
    '_EXE_ALIAS_LIST' : [None, {
        'ListEntry': [0x00, ['_LIST_ENTRY']],
        'ExeLength': [0x08, ['unsigned short']],
        'ExeName': [0x0C, ['pointer', [
            'UnicodeString', dict(
                encoding='utf16',
                length=lambda x: x.ExeLength * 2)]]],
        'AliasList': [0x10, ['_LIST_ENTRY']],
        }],
    '_POPUP_LIST' : [None, {
        'ListEntry' : [0x00, ['_LIST_ENTRY']],
        }],
    '_CONSOLE_INFORMATION': [None, {
        'CurrentScreenBuffer': [0xB0, ['pointer', ['_SCREEN_INFORMATION']]],
        'ScreenBuffer': [0xB4, ['pointer', ['_SCREEN_INFORMATION']]],
        'HistoryList': [0x108, ['_LIST_ENTRY']],
        'ProcessList': [0x100, ['_LIST_ENTRY']],
        'ExeAliasList': [0x110, ['_LIST_ENTRY']],
        'HistoryBufferCount': [0x118, ['unsigned short']],
        'HistoryBufferMax': [0x11A, ['unsigned short']],
        'CommandHistorySize': [0x11C, ['unsigned short']],
        'OriginalTitle': [0x124, ['pointer', [
            'UnicodeString', dict(
                encoding='utf16',
                length=256)]]],
        'Title': [0x128, ['pointer', [
            'UnicodeString', dict(
                encoding='utf16',
                length=256
                )]]],
        }],
    '_CONSOLE_PROCESS': [None, {
        'ListEntry': [0x00, ['_LIST_ENTRY']],
        'ProcessHandle': [0x08, ['unsigned int']],
        'Process': [0x0C, ['pointer', ['_CSR_PROCESS']]],
        }],
    '_SCREEN_INFORMATION': [None, {
        'Console': [0x00, ['pointer', ['_CONSOLE_INFORMATION']]],
        'ScreenX': [0x24, ['short']],
        'ScreenY': [0x26, ['short']],
        'Rows': [0x58, ['pointer', [
            'array', lambda x: x.ScreenY, ['_ROW']]]],
        'Next': [0xF8, ['pointer', ['_SCREEN_INFORMATION']]],
        }],
    '_ROW': [0x1C, {
        'Chars': [0x08, ['pointer', [
            'UnicodeString', dict(
                encoding='utf16', length=256
                )]]],
        }],

    # this is a public PDB
    '_CSR_PROCESS' : [0x60, {
        'ClientId' : [0x0, ['_CLIENT_ID']],
        'ListLink' : [0x8, ['_LIST_ENTRY']],
        'ThreadList' : [0x10, ['_LIST_ENTRY']],
        'NtSession' : [0x18, ['pointer', ['_CSR_NT_SESSION']]],
        'ClientPort' : [0x1c, ['pointer', ['void']]],
        'ClientViewBase' : [0x20, ['pointer', ['unsigned char']]],
        'ClientViewBounds' : [0x24, ['pointer', ['unsigned char']]],
        'ProcessHandle' : [0x28, ['pointer', ['void']]],
        'SequenceNumber' : [0x2c, ['unsigned long']],
        'Flags' : [0x30, ['unsigned long']],
        'DebugFlags' : [0x34, ['unsigned long']],
        'ReferenceCount' : [0x38, ['unsigned long']],
        'ProcessGroupId' : [0x3c, ['unsigned long']],
        'ProcessGroupSequence' : [0x40, ['unsigned long']],
        'LastMessageSequence' : [0x44, ['unsigned long']],
        'NumOutstandingMessages' : [0x48, ['unsigned long']],
        'ShutdownLevel' : [0x4c, ['unsigned long']],
        'ShutdownFlags' : [0x50, ['unsigned long']],
        'Luid' : [0x54, ['_LUID']],
        'ServerDllPerProcessData' : [0x5c, [
            'array', 1, ['pointer', ['void']]]],
        }],
    }

winsrv_types_x64 = {
    '_COMMAND': [None, {
        'CmdLength': [0x00, ['unsigned short']],
        'Cmd' : [0x02, ['UnicodeString', dict(
            encoding='utf16',
            length=lambda x: x.CmdLength)]],
        }],
    '_COMMAND_HISTORY': [None, {
        'Flags' : [0x00, ['Flags', {
            'bitmap': {'Allocated': 0, 'Reset': 1}}]],
        'ListEntry': [0x08, ['_LIST_ENTRY']],
        'Application': [0x18, ['pointer', [
            'UnicodeString', dict(
                encoding='utf16',
                length=256)]]],
        'CommandCount': [0x20, ['short']],
        'LastAdded': [0x22, ['short']],
        'LastDisplayed': [0x24, ['short']],
        'FirstCommand': [0x26, ['short']],
        'CommandCountMax': [0x28, ['short']],
        'ProcessHandle': [0x30, ['unsigned int']],
        'PopupList': [0x38, ['_LIST_ENTRY']],
        'CommandBucket': [0x48, [
            'array', lambda x: x.CommandCount, [
                'pointer', ['_COMMAND']]]],
        }],
    '_ALIAS': [None, {
        'ListEntry': [0x00, ['_LIST_ENTRY']],
        'SourceLength': [0x10, ['unsigned short']],
        'TargetLength': [0x12, ['unsigned short']],
        'Source': [0x14, ['pointer', [
            'UnicodeString', dict(
                encoding='utf16',
                length=lambda x: x.SourceLength * 2)]]],
        'Target': [0x1C, ['pointer', ['UnicodeString', dict(
            encoding='utf16',
            length=lambda x: x.TargetLength * 2)]]],
        }],
    '_EXE_ALIAS_LIST' : [None, {
        'ListEntry': [0x00, ['_LIST_ENTRY']],
        'ExeLength': [0x10, ['unsigned short']],
        'ExeName': [0x12, ['pointer', [
            'UnicodeString', dict(
                encoding='utf16',
                length=lambda x: x.ExeLength * 2)]]],
        'AliasList': [0x1A, ['_LIST_ENTRY']],
        }],
    '_POPUP_LIST' : [None, {
        'ListEntry' : [0x00, ['_LIST_ENTRY']],
        }],
    '_CONSOLE_INFORMATION': [None, {
        'CurrentScreenBuffer': [0xE8, ['pointer', ['_SCREEN_INFORMATION']]],
        'ScreenBuffer': [0xF0, ['pointer', ['_SCREEN_INFORMATION']]],
        'HistoryList': [0x188, ['_LIST_ENTRY']],
        'ProcessList': [0x178, ['_LIST_ENTRY']],
        'ExeAliasList': [0x198, ['_LIST_ENTRY']],
        'HistoryBufferCount': [0x1A8, ['unsigned short']],
        'HistoryBufferMax': [0x1AA, ['unsigned short']],
        'CommandHistorySize': [0x1AC, ['unsigned short']],
        'OriginalTitle': [0x1B0, ['pointer', [
            'UnicodeString', dict(
                encoding='utf16', length=256)]]],
        'Title': [0x1B8, ['pointer', [
            'UnicodeString', dict(
                encoding='utf16', length=256)]]],
        }],
    '_CONSOLE_PROCESS': [None, {
        'ListEntry': [0x00, ['_LIST_ENTRY']],
        'ProcessHandle': [0x10, ['unsigned int']],
        'Process': [0x18, ['pointer', ['_CSR_PROCESS']]],
        }],
    '_SCREEN_INFORMATION': [None, {
        'Console': [0x00, ['pointer', ['_CONSOLE_INFORMATION']]],
        'ScreenX': [0x28, ['short']],
        'ScreenY': [0x2A, ['short']],
        'Rows': [0x68, ['pointer', [
            'array', lambda x: x.ScreenY, ['_ROW']]]],
        'Next': [0x128, ['pointer', ['_SCREEN_INFORMATION']]],
        }],
    '_ROW': [0x28, {
        'Chars': [0x08, ['pointer', [
            'UnicodeString', dict(
                encoding='utf16',
                length=256)]]],
        }],

    # this is a public PDB
    '_CSR_PROCESS' : [0x60, {
        'ClientId' : [0x0, ['_CLIENT_ID']],
        'ListLink' : [0x8, ['_LIST_ENTRY']],
        'ThreadList' : [0x10, ['_LIST_ENTRY']],
        'NtSession' : [0x18, ['pointer', ['_CSR_NT_SESSION']]],
        'ClientPort' : [0x1c, ['pointer', ['void']]],
        'ClientViewBase' : [0x20, ['pointer', ['unsigned char']]],
        'ClientViewBounds' : [0x24, ['pointer', ['unsigned char']]],
        'ProcessHandle' : [0x28, ['pointer', ['void']]],
        'SequenceNumber' : [0x2c, ['unsigned long']],
        'Flags' : [0x30, ['unsigned long']],
        'DebugFlags' : [0x34, ['unsigned long']],
        'ReferenceCount' : [0x38, ['unsigned long']],
        'ProcessGroupId' : [0x3c, ['unsigned long']],
        'ProcessGroupSequence' : [0x40, ['unsigned long']],
        'LastMessageSequence' : [0x44, ['unsigned long']],
        'NumOutstandingMessages' : [0x48, ['unsigned long']],
        'ShutdownLevel' : [0x4c, ['unsigned long']],
        'ShutdownFlags' : [0x50, ['unsigned long']],
        'Luid' : [0x54, ['_LUID']],
        'ServerDllPerProcessData' : [0x5c, [
            'array', 1, ['pointer', ['void']]]],
        }],
    }

class _CONSOLE_INFORMATION(obj.Struct):
    """ object class for console information structs """

    def get_screens(self):
        """Generator for screens in the console.

        A console can have multiple screen buffers at a time,
        but only the current/active one is displayed.

        Multiple screens are tracked using the singly-linked
        list _SCREEN_INFORMATION.Next.

        See CreateConsoleScreenBuffer
        """
        screens = [self.CurrentScreenBuffer]

        if self.ScreenBuffer not in screens:
            screens.append(self.ScreenBuffer)

        for screen in screens:
            cur = screen
            while cur and cur.v() != 0:
                yield cur
                cur = cur.Next.dereference()

class _SCREEN_INFORMATION(obj.Struct):
    """ object class for screen information """

    def get_buffer(self, truncate=True):
        """Get the screen buffer.

        The screen buffer is comprised of the screen's Y
        coordinate which tells us the number of rows and
        the X coordinate which tells us the width of each
        row in characters. These together provide all of
        the input and output that users see when the
        console is displayed.

        @param truncate: True if the empty rows at the
        end (i.e. bottom) of the screen buffer should be
        supressed.
        """
        rows = []

        for _, row in enumerate(self.Rows.dereference()):
            if row.Chars.is_valid():
                rows.append(utils.SmartUnicode(
                    row.Chars.dereference())[0:self.ScreenX])

        # To truncate empty rows at the end, walk the list
        # backwards and get the last non-empty row. Use that
        # row index to splice. An "empty" row isn't just ""
        # as one might assume. It is actually ScreenX number
        # of space characters

        if truncate:
            non_empty_index = 0
            for index, row in enumerate(reversed(rows)):
                ## It seems that when the buffer width is greater than 128
                ## characters, its truncated to 128 in memory.
                if row.count(" ") != min(self.ScreenX, 128):
                    non_empty_index = index
                    break
            if non_empty_index == 0:
                rows = []
            else:
                rows = rows[0:len(rows) - non_empty_index]

        return rows


class WinSrv86(basic.Profile32Bits, basic.BasicClasses):
    """A domain specific profile for the xp, 2008."""

    def __init__(self, **kwargs):
        super(WinSrv86, self).__init__(**kwargs)
        self.add_types(common_types)
        self.add_types(winsrv_types_x86)
        self.add_classes({
            '_CONSOLE_INFORMATION': _CONSOLE_INFORMATION,
            '_SCREEN_INFORMATION': _SCREEN_INFORMATION,
            })


class WinSrv64(basic.ProfileLLP64, basic.BasicClasses):
    """A domain specific profile for the xp, 2008."""

    def __init__(self, **kwargs):
        super(WinSrv64, self).__init__(**kwargs)
        self.add_types(common_types_64)
        self.add_types(winsrv_types_x64)
        self.add_classes({
            '_CONSOLE_INFORMATION': _CONSOLE_INFORMATION,
            '_SCREEN_INFORMATION': _SCREEN_INFORMATION,
            })


class ConHost86(basic.Profile32Bits, basic.BasicClasses):
    """A domain specific profile for windows 7."""

    def __init__(self, **kwargs):
        super(ConHost86, self).__init__(**kwargs)
        self.add_types(common_types)
        self.add_types(conhost_types_x86)
        self.add_classes({
            '_CONSOLE_INFORMATION': _CONSOLE_INFORMATION,
            '_SCREEN_INFORMATION': _SCREEN_INFORMATION,
            })


class ConHost64(basic.ProfileLLP64, basic.BasicClasses):
    """A domain specific profile for windows 7."""

    def __init__(self, **kwargs):
        super(ConHost64, self).__init__(**kwargs)
        self.add_types(common_types_64)
        self.add_types(conhost_types_x64)
        self.add_classes({
            '_CONSOLE_INFORMATION': _CONSOLE_INFORMATION,
            '_SCREEN_INFORMATION': _SCREEN_INFORMATION,
            })


class WinHistoryScanner(vadinfo.VadScanner):
    """A vad scanner for command histories.

    The default pattern we search for, as described by Stevens and Casey, is
    "\x32\x00". That's because CommandCountMax is a little-endian unsigned short
    whose default value is 50. However, that value can be changed by right
    clicking cmd.exe and going to Properties->Options->Cmd History or by calling
    the API function kernel32!SetConsoleHistoryInfo. Thus you can tweak the
    search criteria by using the --MAX_HISTORY.
    """
    def __init__(self, max_history=MAX_HISTORY_DEFAULT, **kwargs):
        super(WinHistoryScanner, self).__init__(**kwargs)
        self.max_history = max_history

    def scan(self, **kwargs):
        for hit in super(WinHistoryScanner, self).scan(**kwargs):
            # Check to see if the object is valid.
            hist = self.profile.Object(
                "_COMMAND_HISTORY", vm=self.address_space,
                offset=hit - self.profile.get_obj_offset(
                    "_COMMAND_HISTORY", "CommandCountMax"))

            if not hist.is_valid():
                continue

            # The count must be between zero and max
            if hist.CommandCount < 0 or hist.CommandCount > self.max_history:
                continue

            # Last added must be between -1 and max
            if hist.LastAdded < -1 or hist.LastAdded > self.max_history:
                continue

            # Last displayed must be between -1 and max
            if hist.LastDisplayed < -1 or hist.LastDisplayed > self.max_history:
                continue

            # First command must be between zero and max
            if hist.FirstCommand < 0 or hist.FirstCommand > self.max_history:
                continue

            # Validate first command with last added
            if (hist.FirstCommand != 0 and
                    hist.FirstCommand != hist.LastAdded + 1):
                continue

            # Process handle must be a valid pid
            if hist.ProcessHandle <= 0 or hist.ProcessHandle > 0xFFFF:
                continue

            Popup = self.profile._POPUP_LIST(
                offset=hist.PopupList.Flink, vm=self.address_space)

            # Check that the popup list entry is in tact
            if Popup.ListEntry.Blink != hist.PopupList.obj_offset:
                continue

            yield hist


class CmdScan(common.WindowsCommandPlugin):
    """Extract command history by scanning for _COMMAND_HISTORY"""
    __name = "cmdscan"

    __args = [
        dict(name="max_history", default=MAX_HISTORY_DEFAULT,
             type="IntParser", help="Value of history buffer size. See "
             "HKEY_CURRENT_USER\\Console\\HistoryBufferSize "
             "for default.")
    ]

    def generate_hits(self):
        """Generates _COMMAND_HISTORY objects."""
        architecture = self.profile.metadata("arch")
        # The process we select is conhost on Win7 or csrss for others

        if self.profile.metadata("major") >= 6:
            process_name = "conhost.exe"
            if architecture == "AMD64":
                process_profile = ConHost64(session=self.session)
            else:
                process_profile = ConHost86(session=self.session)
        else:
            process_name = "csrss.exe"
            if architecture == "AMD64":
                process_profile = WinSrv64(session=self.session)
            else:
                process_profile = WinSrv86(session=self.session)

        # Only select those processes we care about:
        for task in self.session.plugins.pslist(
                proc_regex=process_name).filter_processes():
            scanner = WinHistoryScanner(
                task=task, process_profile=process_profile,
                max_history=self.plugin_args.max_history, session=self.session)

            pattern = bytes((self.plugin_args.max_history, 0))
            scanner.checks = [
                ("StringCheck", dict(needle=pattern))
                ]
            scanner.build_constraints()

            for hist in scanner.scan():
                yield task, hist

    def render(self, renderer):
        for task, hist in self.generate_hits():
            renderer.section()
            renderer.format(u"CommandProcess: {0} Pid: {1}\n",
                            task.ImageFileName, task.UniqueProcessId)
            renderer.format(
                u"CommandHistory: {0:#x} Application: {1} Flags: {2}\n",
                hist.obj_offset, hist.Application.dereference(), hist.Flags)

            renderer.format(
                u"CommandCount: {0} LastAdded: {1} LastDisplayed: {2}\n",
                hist.CommandCount, hist.LastAdded, hist.LastDisplayed)

            renderer.format(u"FirstCommand: {0} CommandCountMax: {1}\n",
                            hist.FirstCommand, hist.CommandCountMax)

            renderer.format(u"ProcessHandle: {0:#x}\n", hist.ProcessHandle)


            renderer.table_header([("Cmd", "command", ">3"),
                                   ("Address", "address", "[addrpad]"),
                                   ("Text", "text", "")])

            # If the _COMMAND_HISTORY is in use, we would only take
            # hist.CommandCount but since we're brute forcing, try the
            # maximum and hope that some slots were not overwritten
            # or zero-ed out.
            pointers = hist.obj_profile.Array(
                target="address", count=hist.CommandCountMax,
                offset=hist.obj_offset + hist.obj_profile.get_obj_offset(
                    "_COMMAND_HISTORY", "CommandBucket"),
                vm=hist.obj_vm)

            for i, p in enumerate(pointers):
                cmd = p.cast("Pointer", target="_COMMAND").deref()
                if cmd.obj_offset and cmd.Cmd:
                    renderer.table_row(
                        i, cmd.obj_offset,
                        utils.SmartUnicode(cmd.Cmd).encode("unicode_escape"))


class ConsoleScanner(vadinfo.VadScanner):
    """A scanner for _CONSOLE_INFORMATION."""

    def __init__(self, max_history=MAX_HISTORY_DEFAULT,
                 history_buffers=HISTORY_BUFFERS_DEFAULTS, **kwargs):
        super(ConsoleScanner, self).__init__(**kwargs)
        self.max_history = max_history
        self.history_buffers = history_buffers

    def scan(self, **kwargs):
        for hit in super(ConsoleScanner, self).scan(**kwargs):
            # Check to see if the object is valid.
            console = self.profile.Object(
                "_CONSOLE_INFORMATION", offset=hit -
                self.profile.get_obj_offset(
                    "_CONSOLE_INFORMATION", "CommandHistorySize"),
                vm=self.address_space, parent=self.task)

            if (console.HistoryBufferMax != self.history_buffers or
                console.HistoryBufferCount > self.history_buffers):
                continue

            # Check the first command history as the final constraint
            next_history = console.HistoryList.Flink.dereference(
                ).dereference_as("_COMMAND_HISTORY", "ListEntry")
            if next_history.CommandCountMax != self.max_history:
                continue

            yield console


class ConsoleScan(CmdScan):
    """Extract command history by scanning for _CONSOLE_INFORMATION"""

    __name = "consolescan"

    __args = [
        dict(name="history_buffers", default=HISTORY_BUFFERS_DEFAULTS,
             type="IntParser",
             help="Value of history buffer size. See "
             "HKEY_CURRENT_USER\\Console\\HistoryBufferSize "
             "for default.")
    ]

    def generate_hits(self):
        """Generates _CONSOLE_INFORMATION objects."""
        architecture = self.profile.metadata("arch")
        # The process we select is conhost on Win7 or csrss for others

        # Only select those processes we care about:
        for task in self.session.plugins.pslist(
                proc_regex="(conhost.exe|csrss.exe)").filter_processes():

            if utils.SmartUnicode(task.ImageFileName).lower() == "conhost.exe":
                if architecture == "AMD64":
                    process_profile = ConHost64(session=self.session)
                else:
                    process_profile = ConHost86(session=self.session)

            elif utils.SmartUnicode(task.ImageFileName).lower() == "csrss.exe":
                if architecture == "AMD64":
                    process_profile = WinSrv64(session=self.session)
                else:
                    process_profile = WinSrv86(session=self.session)

            else: continue

            scanner = ConsoleScanner(
                task=task, process_profile=process_profile,
                session=self.session,
                max_history=self.plugin_args.max_history,
                history_buffers=self.plugin_args.history_buffers)

            pattern = bytes((self.plugin_args.max_history, 0))
            scanner.checks = [
                ("StringCheck", dict(needle=pattern))
                ]
            scanner.build_constraints()

            for console in scanner.scan():
                yield task, console

    def render(self, renderer):
        for task, console in self.generate_hits():
            renderer.section()

            renderer.format(u"ConsoleProcess: {0} Pid: {1}\n",
                            task.ImageFileName, task.UniqueProcessId)

            renderer.format(u"Console: {0:#x} CommandHistorySize: {1}\n",
                            console, console.CommandHistorySize)

            renderer.format(
                u"HistoryBufferCount: {0} HistoryBufferMax: {1}\n",
                console.HistoryBufferCount, console.HistoryBufferMax)

            renderer.format(u"OriginalTitle: {0}\n",
                            console.OriginalTitle.dereference())

            renderer.format(u"Title: {0}\n", console.Title.dereference())

            for console_proc in console.ProcessList.list_of_type(
                    "_CONSOLE_PROCESS", "ListEntry"):
                process = task.ObReferenceObjectByHandle(
                    console_proc.ProcessHandle, type="_EPROCESS")

                if process:
                    renderer.format(
                        u"AttachedProcess: {0} Pid: {1} Handle: {2:#x}\n",
                        process.ImageFileName, process.UniqueProcessId,
                        console_proc.ProcessHandle)

            for hist in console.HistoryList.list_of_type(
                    "_COMMAND_HISTORY", "ListEntry"):
                renderer.format(u"----\n")

                renderer.format(
                    u"CommandHistory: {0:#x} Application: {1} Flags: {2}\n",
                    hist, hist.Application.dereference(),
                    hist.Flags)
                renderer.format(
                    u"CommandCount: {0} LastAdded: {1} LastDisplayed: {2}\n",
                    hist.CommandCount, hist.LastAdded, hist.LastDisplayed)

                renderer.format(u"FirstCommand: {0} CommandCountMax: {1}\n",
                                hist.FirstCommand, hist.CommandCountMax)

                renderer.format(u"ProcessHandle: {0:#x}\n", hist.ProcessHandle)
                for i, cmd in enumerate(hist.CommandBucket):
                    if cmd.Cmd.is_valid():
                        renderer.format(u"Cmd #{0} at {1:#x}: {2}\n",
                                        i, cmd, str(cmd.Cmd))

            for exe_alias in console.ExeAliasList.list_of_type(
                    "_EXE_ALIAS_LIST", "ListEntry"):

                for alias in exe_alias.AliasList.list_of_type(
                        "_ALIAS", "ListEntry"):
                    renderer.format(u"----\n")
                    renderer.format(
                        u"Alias: {0} Source: {1} Target: {2}\n",
                        exe_alias.ExeName.dereference(),
                        alias.Source.dereference(),
                        alias.Target.dereference())

            for screen in console.get_screens():
                renderer.format(u"----\n")
                renderer.format(u"Screen {0:#x} X:{1} Y:{2}\n",
                                screen.dereference(), screen.ScreenX,
                                screen.ScreenY)

                renderer.format(u"Dump:\n{0}\n", '\n'.join(screen.get_buffer()))


class Conhost(pe_vtypes.BasicPEProfile):
    """A profile for Conhost.exe."""

    @classmethod
    def Initialize(cls, profile):
        super(Conhost, cls).Initialize(profile)
        if profile.metadata("arch") == "AMD64":
            profile.add_overlay(conhost_types_x64)
            profile.add_overlay(common_types_64)

        elif profile.metadata("arch") == "I386":
            profile.add_overlay(conhost_types_x86)
            profile.add_overlay(common_types)


class Consoles(common.WindowsCommandPlugin):
    """Enumerate command consoles."""

    name = "consoles"

    def _render_conhost_process(self, task, renderer):
        self.cc.SwitchProcessContext(task)
        console = self.session.address_resolver.get_constant_object(
            "conhost!gConsoleInformation", "_CONSOLE_INFORMATION")

        if console == None:
            self.session.logging.warning(
                "Unable to load profile for conhost.exe.")
            return

        renderer.format(u"ConsoleProcess: {0}\n", task)
        renderer.format(u"Title: {0}\n", console.Title.deref())
        renderer.format(u"OriginalTitle: {0}\n",
                        console.OriginalTitle.dereference())

        for console_proc in console.ProcessList.list_of_type(
                "_CONSOLE_PROCESS", "ListEntry"):
            process = task.ObReferenceObjectByHandle(
                console_proc.ProcessHandle, type="_EPROCESS")

            if process:
                renderer.format(
                    u"AttachedProcess: {0} Handle: {1:style=address}\n",
                    process, console_proc.ProcessHandle)

        for hist in console.HistoryList.list_of_type(
                "_COMMAND_HISTORY", "ListEntry"):
            renderer.format(u"----\n")

            renderer.format(
                u"CommandHistory: {0:#x} Application: {1} Flags: {2}\n",
                hist, hist.Application.dereference(),
                hist.Flags)

            renderer.format(
                u"CommandCount: {0} LastAdded: {1} LastDisplayed: {2}\n",
                hist.CommandCount, hist.LastAdded, hist.LastDisplayed)

            renderer.format(u"FirstCommand: {0} CommandCountMax: {1}\n",
                            hist.FirstCommand, hist.CommandCountMax)

            renderer.format(u"ProcessHandle: {0:#x}\n", hist.ProcessHandle)
            for i, cmd in enumerate(hist.CommandBucket):
                if cmd.Cmd.is_valid():
                    renderer.format(u"Cmd #{0} at {1:#x}: {2}\n",
                                    i, cmd, str(cmd.Cmd))

        for exe_alias in console.ExeAliasList.list_of_type(
                "_EXE_ALIAS_LIST", "ListEntry"):

            for alias in exe_alias.AliasList.list_of_type(
                    "_ALIAS", "ListEntry"):
                if alias == None:
                    continue

                renderer.format(u"----\n")
                renderer.format(
                    u"Alias: {0} Source: {1} Target: {2}\n",
                    exe_alias.ExeName.dereference(),
                    alias.Source.dereference(),
                    alias.Target.dereference())

        for screen in console.ScreenBuffer.walk_list("Next"):
            renderer.format(u"----\n")
            renderer.format(u"Screen {0:#x} X:{1} Y:{2}\n",
                            screen, screen.ScreenX, screen.ScreenY)

            screen_data = []
            for row in screen.Rows:
                row_data = row.Chars.deref()

                # Since a row is always screen.ScreenX wide a 0 width row means
                # the page is not available or incorrect.
                if len(row_data) == 0:
                    continue

                screen_data.append((row_data, str(row_data).rstrip()))

            # Trim empty rows from the end of the screen.
            while screen_data and not screen_data[-1][1]:
                screen_data.pop(-1)

            for row, row_string in screen_data:
                renderer.format(
                    "{0:style=address}: {1}\n", row, row_string)

    def render(self, renderer):
        self.cc = self.session.plugins.cc()
        with self.cc:
            for task in self.session.plugins.pslist(
                    proc_regex="conhost.exe").filter_processes():
                renderer.section()
                self._render_conhost_process(task, renderer)
