# Andrew Davis, andrew.davis@mandiant.com
# Copyright 2012 Mandiant
#
# Mandiant licenses this file to you under the Apache License, Version
# 2.0 (the "License"); you may not use this file except in compliance with the
# License.  You may obtain a copy of the License at:
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.  See the License for the specific language governing
# permissions and limitations under the License.
#
# Identifies and parses Application Compatibility Shim Cache entries for forensic data.


# Original Mandiant shim cache parser, ported to Python 3
# Identifies and parses Application Compatibility Shim Cache entries for forensic data.

import datetime
import io as sio
import struct

import logging
import pytz

CACHE_MAGIC_NT5_2 = 0xBADC0FFE
CACHE_HEADER_SIZE_NT5_2 = 0x8
NT5_2_ENTRY_SIZE32 = 0x18
NT5_2_ENTRY_SIZE64 = 0x20

# Values used by Windows 6.1 (Win7 and Server 2008 R2)
CACHE_MAGIC_NT6_1 = 0xBADC0FEE
CACHE_HEADER_SIZE_NT6_1 = 0x80
NT6_1_ENTRY_SIZE32 = 0x20
NT6_1_ENTRY_SIZE64 = 0x30
CSRSS_FLAG = 0x2

# Values used by Windows 5.1 (WinXP 32-bit)
WINXP_MAGIC32 = 0xDEADBEEF
WINXP_HEADER_SIZE32 = 0x190
WINXP_ENTRY_SIZE32 = 0x228
MAX_PATH = 520

# Values used by Windows 8
WIN8_STATS_SIZE = 0x80
WIN8_MAGIC = b"00ts"

# Magic value used by Windows 8.1
WIN81_MAGIC = b"10ts"

# Values used by Windows 10
WIN10_STATS_SIZE = 0x30
WIN10_MAGIC = b"10ts"
CACHE_HEADER_SIZE_NT6_4 = 0x30
CACHE_MAGIC_NT6_4 = 0x30

BAD_ENTRY_DATA = "N/A"
G_VERBOSE = False
G_USE_BOM = False
OUTPUT_HEADER = ["Last Modified", "Last Update", "Path", "File Size", "Exec Flag"]

logger = logging.getLogger(__name__)


# Shim Cache format used by Windows 5.2 and 6.0 (Server 2003 through Vista/Server 2008)
class CacheEntryNt5(object):
    def __init__(self, is_32_bit, data=None):
        self.w_length = None
        self.w_maximum_length = None
        self.offset = None
        self.dw_low_date_time = None
        self.dw_high_date_time = None
        self.dw_file_size_low = None
        self.dw_file_size_high = None
        self.is_32_bit = is_32_bit
        if data is not None:
            self.update(data)

    def update(self, data):
        if self.is_32_bit:
            entry = struct.unpack("<2H 3L 2L", data)
        else:
            entry = struct.unpack("<2H 4x Q 2L 2L", data)
        self.w_length = entry[0]
        self.w_maximum_length = entry[1]
        self.offset = entry[2]
        self.dw_low_date_time = entry[3]
        self.dw_high_date_time = entry[4]
        self.dw_file_size_low = entry[5]
        self.dw_file_size_high = entry[6]

    def size(self):

        if self.is_32_bit:
            return NT5_2_ENTRY_SIZE32
        else:
            return NT5_2_ENTRY_SIZE64


# Shim Cache format used by Windows 6.1 (Win7 through Server 2008 R2).
class CacheEntryNt6(object):
    def __init__(self, is_32_bit, data=None):
        self.w_length = None
        self.w_maximum_length = None
        self.offset = None
        self.dw_low_date_time = None
        self.dw_high_date_time = None
        self.file_flags = None
        self.flags = None
        self.blob_size = None
        self.blob_offset = None
        self.is_32_bit = is_32_bit
        if data is not None:
            self.update(data)

    def update(self, data):
        if self.is_32_bit:
            entry = struct.unpack("<2H 7L", data)
        else:
            entry = struct.unpack("<2H 4x Q 4L 2Q", data)
        self.w_length = entry[0]
        self.w_maximum_length = entry[1]
        self.offset = entry[2]
        self.dw_low_date_time = entry[3]
        self.dw_high_date_time = entry[4]
        self.file_flags = entry[5]
        self.flags = entry[6]
        self.blob_size = entry[7]
        self.blob_offset = entry[8]

    def size(self):
        if self.is_32_bit:
            return NT6_1_ENTRY_SIZE32
        else:
            return NT6_1_ENTRY_SIZE64


# Convert FILETIME to datetime.
# Based on http://code.activestate.com/recipes/511425-filetime-to-datetime/
def convert_filetime(dw_low_date_time, dw_high_date_time):
    try:
        date = datetime.datetime(1601, 1, 1, 0, 0, 0)
        temp_time = dw_high_date_time
        temp_time <<= 32
        temp_time |= dw_low_date_time
        return pytz.utc.localize(date + datetime.timedelta(microseconds=temp_time / 10))
    except OverflowError:
        return None


# Return a unique list while preserving ordering.
def unique_list(li):
    ret_list = []
    for entry in li:
        if entry not in ret_list:
            ret_list.append(entry)
    return ret_list


# Read the Shim Cache format, return a list of last modified dates/paths.
def get_shimcache_entries(cachebin, as_json=False):

    if len(cachebin) < 16:
        # Data size less than minimum header size.
        return None

    # Get the format type
    magic = struct.unpack("<L", cachebin[0:4])[0]

    # This is a Windows 2k3/Vista/2k8 Shim Cache format,
    if magic == CACHE_MAGIC_NT5_2:

        # Shim Cache types can come in 32-bit or 64-bit formats. We can
        # determine this because 64-bit entries are serialized with u_int64
        # pointers. This means that in a 64-bit entry, valid UNICODE_STRING
        # sizes are followed by a NULL DWORD. Check for this here.
        test_size = struct.unpack("<H", cachebin[8:10])[0]
        test_max_size = struct.unpack("<H", cachebin[10:12])[0]
        if (
            test_max_size - test_size == 2 and struct.unpack("<L", cachebin[12:16])[0]
        ) == 0:
            logger.debug("[+] Found 64bit Windows 2k3/Vista/2k8 Shim Cache data...")
            entry = CacheEntryNt5(False)
            yield from read_nt5_entries(cachebin, entry, as_json=as_json)

        # Otherwise it's 32-bit data.
        else:
            logger.debug("[+] Found 32bit Windows 2k3/Vista/2k8 Shim Cache data...")
            entry = CacheEntryNt5(True)
            yield from read_nt5_entries(cachebin, entry, as_json=as_json)

    # This is a Windows 7/2k8-R2 Shim Cache.
    elif magic == CACHE_MAGIC_NT6_1:
        test_size = struct.unpack(
            "<H", cachebin[CACHE_HEADER_SIZE_NT6_1: CACHE_HEADER_SIZE_NT6_1 + 2]
        )[0]
        test_max_size = struct.unpack(
            "<H", cachebin[CACHE_HEADER_SIZE_NT6_1 + 2: CACHE_HEADER_SIZE_NT6_1 + 4]
        )[0]

        # Shim Cache types can come in 32-bit or 64-bit formats.
        # We can determine this because 64-bit entries are serialized with
        # u_int64 pointers. This means that in a 64-bit entry, valid
        # UNICODE_STRING sizes are followed by a NULL DWORD. Check for this here.
        if (
            test_max_size - test_size == 2
            and struct.unpack(
                "<L",
                cachebin[CACHE_HEADER_SIZE_NT6_1 + 4: CACHE_HEADER_SIZE_NT6_1 + 8],
            )[0]
        ) == 0:
            logger.debug("[+] Found 64bit Windows 7/2k8-R2 Shim Cache data...")
            entry = CacheEntryNt6(False)
            yield from read_nt6_entries(cachebin, entry, as_json=as_json)
        else:
            logger.debug("[+] Found 32bit Windows 7/2k8-R2 Shim Cache data...")
            entry = CacheEntryNt6(True)
            yield from read_nt6_entries(cachebin, entry, as_json=as_json)

    # This is WinXP cache data
    elif magic == WINXP_MAGIC32:
        logger.debug("[+] Found 32bit Windows XP Shim Cache data...")
        yield from read_winxp_entries(cachebin, as_json=as_json)

    # Check the data set to see if it matches the Windows 8 format.
    elif (
        len(cachebin) > WIN8_STATS_SIZE
        and cachebin[WIN8_STATS_SIZE: WIN8_STATS_SIZE + 4] == WIN8_MAGIC
    ):
        logger.debug("[+] Found Windows 8/2k12 Apphelp Cache data...")
        yield from read_win8_entries(cachebin, WIN8_MAGIC, as_json=as_json)

    # Windows 8.1 will use a different magic dword, check for it
    elif (
        len(cachebin) > WIN8_STATS_SIZE
        and cachebin[WIN8_STATS_SIZE: WIN8_STATS_SIZE + 4] == WIN81_MAGIC
    ):
        logger.debug("[+] Found Windows 8.1 Apphelp Cache data...")
        yield from read_win8_entries(cachebin, WIN81_MAGIC, as_json=as_json)

    # Windows 10 will use a different magic dword, check for it
    elif (
        len(cachebin) > WIN10_STATS_SIZE
        and cachebin[WIN10_STATS_SIZE: WIN10_STATS_SIZE + 4] == WIN10_MAGIC
    ):
        logger.debug("[+] Found Windows 10 Apphelp Cache data...")
        yield from read_win10_entries(cachebin, WIN10_MAGIC, as_json=as_json)

    # Windows 10 creators update moved the damn magic 4 bytes forward...
    elif (
        len(cachebin) > WIN10_STATS_SIZE
        and cachebin[WIN10_STATS_SIZE + 4: WIN10_STATS_SIZE + 8] == WIN10_MAGIC
    ):
        logger.debug("[+] Found Windows 10 Apphelp Cache data... (creators update)")
        yield from read_win10_entries(
            cachebin, WIN10_MAGIC, creators_update=True, as_json=as_json
        )

    else:
        raise Exception(
            "Got an unrecognized magic value of 0x{:x}... bailing".format(magic)
        )


# Read Windows 8/2k12/8.1 Apphelp Cache entry formats.
def read_win8_entries(bin_data, ver_magic, as_json=False):

    entry_meta_len = 12

    # Skip past the stats in the header
    cache_data = bin_data[WIN8_STATS_SIZE:]

    data = sio.BytesIO(cache_data)
    while data.tell() < len(cache_data):
        header = data.read(entry_meta_len)
        # Read in the entry metadata
        # Note: the crc32 hash is of the cache entry data
        magic, crc32_hash, entry_len = struct.unpack("<4sLL", header)

        # Check the magic tag
        if magic != ver_magic:
            raise Exception(
                "Invalid version magic tag found: {}".format(
                    struct.unpack("I", magic)[0]
                )
            )

        entry_data = sio.BytesIO(data.read(entry_len))

        # Read the path length
        path_len = struct.unpack("<H", entry_data.read(2))[0]
        if path_len == 0:
            path = "None"
        else:
            path = entry_data.read(path_len).decode("utf-16le", "replace")

        # Check for package data
        package_len = struct.unpack("<H", entry_data.read(2))[0]
        if package_len > 0:
            # Just skip past the package data if present (for now)
            entry_data.seek(package_len, 1)

        # Read the remaining entry data
        flags, unk_1, low_datetime, high_datetime, unk_2 = struct.unpack(
            "<LLLLL", entry_data.read(20)
        )

        # Check the flag set in CSRSS
        if flags & CSRSS_FLAG:
            exec_flag = "True"
        else:
            exec_flag = "False"

        last_mod_date = convert_filetime(low_datetime, high_datetime)

        yield {
            "last_mod_date": last_mod_date.isoformat() if as_json else last_mod_date,
            "path": path,
            "exec_flag": exec_flag,
        }


# Read Windows 10 Apphelp Cache entry format
def read_win10_entries(bin_data, ver_magic, creators_update=False, as_json=False):

    entry_meta_len = 12

    # Skip past the stats in the header
    if creators_update:
        cache_data = bin_data[WIN10_STATS_SIZE + 4:]
    else:
        cache_data = bin_data[WIN10_STATS_SIZE:]

    data = sio.BytesIO(cache_data)
    while data.tell() < len(cache_data):
        header = data.read(entry_meta_len)
        # Read in the entry metadata
        # Note: the crc32 hash is of the cache entry data
        magic, crc32_hash, entry_len = struct.unpack("<4sLL", header)

        # Check the magic tag
        if magic != ver_magic:
            raise Exception(
                "Invalid version magic tag found: {}".format(
                    struct.unpack("<I", magic)[0]
                )
            )

        entry_data = sio.BytesIO(data.read(entry_len))

        # Read the path length
        path_len = struct.unpack("<H", entry_data.read(2))[0]
        if path_len == 0:
            path = "None"
        else:
            path = entry_data.read(path_len).decode("utf-16le", "replace")

        # Read the remaining entry data
        low_datetime, high_datetime = struct.unpack("<LL", entry_data.read(8))

        # Skip the unrecognized Microsoft App entry format for now
        if not (low_datetime + high_datetime):
            continue
        else:
            last_mod_date = convert_filetime(low_datetime, high_datetime)

        yield {
            "last_mod_date": last_mod_date.isoformat() if as_json else last_mod_date,
            "path": path,
        }


# Read Windows 2k3/Vista/2k8 Shim Cache entry formats.
def read_nt5_entries(bin_data, entry, as_json=False):
    contains_file_size = False
    entry_size = entry.size()

    num_entries = struct.unpack("<L", bin_data[4:8])[0]
    if num_entries == 0:
        return None

    # On Windows Server 2008/Vista, the filesize is swapped out of this
    # structure with two 4-byte flags. Check to see if any of the values in
    # "dwFileSizeLow" are larger than 2-bits. This indicates the entry contained file sizes.
    for offset in range(
        CACHE_HEADER_SIZE_NT5_2,
        (num_entries * entry_size) + CACHE_HEADER_SIZE_NT5_2,
        entry_size,
    ):

        entry.update(bin_data[offset: offset + entry_size])

        if entry.dw_file_size_low > 3:
            contains_file_size = True
            break

    # Now grab all the data in the value.
    for offset in range(
        CACHE_HEADER_SIZE_NT5_2,
        (num_entries * entry_size) + CACHE_HEADER_SIZE_NT5_2,
        entry_size,
    ):

        entry.update(bin_data[offset: offset + entry_size])

        last_mod_date = convert_filetime(
            entry.dw_low_date_time, entry.dw_high_date_time
        )
        path = bin_data[entry.offsets: entry.offset + entry.w_length].decode(
            "utf-16le", "replace"
        )

        # It contains file size data.
        exec_flag = None
        if contains_file_size:
            yield {
                "last_mod_date": (
                    last_mod_date.isoformat() if as_json else last_mod_date
                ),
                "path": path,
                "file_size": entry.dw_file_size_low,
            }

        # It contains flags.
        else:
            # Check the flag set in CSRSS
            if entry.dw_file_size_low & CSRSS_FLAG:
                exec_flag = "True"
            else:
                exec_flag = "False"

            yield {
                "last_mod_date": (
                    last_mod_date.isoformat() if as_json else last_mod_date
                ),
                "path": path,
                "exec_flag": exec_flag,
            }


# Read the Shim Cache Windows 7/2k8-R2 entry format,
# return a list of last modifed dates/paths.
def read_nt6_entries(bin_data, entry, as_json=False):
    entry_size = entry.size()
    num_entries = struct.unpack("<L", bin_data[4:8])[0]

    if num_entries == 0:
        return None

    # Walk each entry in the data structure.
    for offset in range(
        CACHE_HEADER_SIZE_NT6_1,
        num_entries * entry_size + CACHE_HEADER_SIZE_NT6_1,
        entry_size,
    ):

        entry.update(bin_data[offset: offset + entry_size])
        last_mod_date = convert_filetime(
            entry.dw_low_date_time, entry.dw_high_date_time
        )
        path = bin_data[entry.offset: entry.offset + entry.w_length].decode(
            "utf-16le", "replace"
        )

        # Test to see if the file may have been executed.
        if entry.file_flags & CSRSS_FLAG:
            exec_flag = "True"
        else:
            exec_flag = "False"

        yield {
            "last_mod_date": last_mod_date.isoformat() if as_json else last_mod_date,
            "path": path,
            "exec_flag": exec_flag,
        }


# Read the WinXP Shim Cache data. Some entries can be missing data but still
# contain useful information, so try to get as much as we can.
def read_winxp_entries(bin_data, as_json=False):

    num_entries = struct.unpack("<L", bin_data[8:12])[0]
    if num_entries == 0:
        return None

    for offset in range(
        WINXP_HEADER_SIZE32,
        (num_entries * WINXP_ENTRY_SIZE32) + WINXP_HEADER_SIZE32,
        WINXP_ENTRY_SIZE32,
    ):
        # No size values are included in these entries, so search for utf-16 terminator.
        path_len = bin_data[offset: offset + (MAX_PATH + 8)].find(b"\x00\x00")

        # if path is corrupt, procede to next entry.
        if path_len == 0:
            continue
        path = bin_data[offset: offset + path_len + 1].decode("utf-16le")

        if len(path) == 0:
            continue

        entry_data = offset + (MAX_PATH + 8)

        # Get last mod time.
        last_mod_time = struct.unpack("<2L", bin_data[entry_data: entry_data + 8])
        last_mod_time = convert_filetime(last_mod_time[0], last_mod_time[1])

        # Get last file size.
        file_size = struct.unpack("<2L", bin_data[entry_data + 8: entry_data + 16])[0]
        if file_size == 0:
            file_size = BAD_ENTRY_DATA

        # Get last update time.
        exec_time = struct.unpack("<2L", bin_data[entry_data + 16: entry_data + 24])
        exec_time = convert_filetime(exec_time[0], exec_time[1])

        yield {
            "last_mod_time": last_mod_time.isoformat() if as_json else last_mod_time,
            "exec_time": exec_time.isoformat() if as_json else exec_time,
            "path": path,
            "file_size": file_size,
        }


def parse_output(output):
    new_output_list = list()
    for row in output:
        exec_flag = False
        if row[4] == "True":
            exec_flag = True
        entry = {
            "last_mod_date": row[0],
            "last_update": row[1],
            "path": row[2],
            "file_size": row[3],
            "exec_flag": exec_flag,
        }
        new_output_list.append(entry)
    return new_output_list
