Reference CIT parser

cit.py

#!/usr/bin/env python3
import array
import argparse
import io
import struct
import sys
from binascii import crc32
from datetime import datetime, timedelta, timezone

try:
    from dissect import cstruct
    from Registry import Registry
except ImportError:
    print("Missing dependencies, run:\npip install dissect.cstruct python-registry")
    sys.exit(1)


try:
    from zoneinfo import ZoneInfo

    HAS_ZONEINFO = True
except ImportError:
    HAS_ZONEINFO = False


cit_def = """
typedef QWORD   FILETIME;

flag TELEMETRY_ANSWERS {
    POWERBROADCAST  = 0x10000,
    DEVICECHANGE    = 0x20000,
    IME_CONTROL     = 0x40000,
    WINHELP         = 0x80000,
};

typedef struct _CIT_HEADER {
    WORD    MajorVersion;
    WORD    MinorVersion;
    DWORD   Size;                   /* Size of the entire buffer */
    FILETIME    CurrentTimeLocal;   /* Maybe the time when the saved CIT was last updated? */
    DWORD   Crc32;                  /* Crc32 of the entire buffer, skipping this field */
    DWORD   EntrySize;
    DWORD   EntryCount;
    DWORD   EntryDataOffset;
    DWORD   SystemDataSize;
    DWORD   SystemDataOffset;
    DWORD   BaseUseDataSize;
    DWORD   BaseUseDataOffset;
    FILETIME    StartTimeLocal;     /* Presumably when the aggregation started */
    FILETIME    PeriodStartLocal;   /* Presumably the starting point of the aggregation period */
    DWORD   AggregationPeriodInS;   /* Presumably the duration over which this data was gathered
                                     * Always 604800 (7 days) */
    DWORD   BitPeriodInS;           /* Presumably the amount of seconds a single bit represents
                                     * Always 3600 (1 hour) */
    DWORD   SingleBitmapSize;       /* This appears to be the sizes of the Stats buffers, always 21 */
    DWORD   _Unk0;                  /* Always 0x00000100? */
    DWORD   HeaderSize;
    DWORD   _Unk1;                  /* Always 0x00000000? */
} CIT_HEADER;

typedef struct _CIT_PERSISTED {
    DWORD   BitmapsOffset;          /* Array of Offset and Size (DWORD, DWORD) */
    DWORD   BitmapsSize;
    DWORD   SpanStatsOffset;        /* Array of Count and Duration (DWORD, DWORD) */
    DWORD   SpanStatsSize;
    DWORD   StatsOffset;            /* Array of WORD */
    DWORD   StatsSize;
} CIT_PERSISTED;

typedef struct _CIT_ENTRY {
    DWORD   ProgramDataOffset;      /* Offset to CIT_PROGRAM_DATA */
    DWORD   UseDataOffset;          /* Offset to CIT_PERSISTED */
    DWORD   ProgramDataSize;
    DWORD   UseDataSize;
} CIT_ENTRY;

typedef struct _CIT_PROGRAM_DATA {
    DWORD   FilePathOffset;         /* Offset to UTF-16-LE file path string */
    DWORD   FilePathSize;           /* strlen of string */
    DWORD   CommandLineOffset;      /* Offset to UTF-16-LE command line string */
    DWORD   CommandLineSize;        /* strlen of string */
    DWORD   PeTimeDateStamp;        /* aka Extra1 */
    DWORD   PeCheckSum;             /* aka Extra2 */
    DWORD   Extra3;                 /* aka Extra3, some flag from PROCESSINFO struct */
} CIT_PROGRAM_DATA;

typedef struct _CIT_BITMAP_ITEM {
    DWORD   Offset;
    DWORD   Size;
} CIT_BITMAP_ITEM;

typedef struct _CIT_SPAN_STAT_ITEM {
    DWORD   Count;
    DWORD   Duration;
} CIT_SPAN_STAT_ITEM;

typedef struct _CIT_SYSTEM_DATA_SPAN_STATS {
    CIT_SPAN_STAT_ITEM  ContextFlushes0;
    CIT_SPAN_STAT_ITEM  Foreground0;
    CIT_SPAN_STAT_ITEM  Foreground1;
    CIT_SPAN_STAT_ITEM  DisplayPower0;
    CIT_SPAN_STAT_ITEM  DisplayRequestChange;
    CIT_SPAN_STAT_ITEM  DisplayPower1;
    CIT_SPAN_STAT_ITEM  DisplayPower2;
    CIT_SPAN_STAT_ITEM  DisplayPower3;
    CIT_SPAN_STAT_ITEM  ContextFlushes1;
    CIT_SPAN_STAT_ITEM  Foreground2;
    CIT_SPAN_STAT_ITEM  ContextFlushes2;
} CIT_SYSTEM_DATA_SPAN_STATS;

typedef struct _CIT_USE_DATA_SPAN_STATS {
    CIT_SPAN_STAT_ITEM  ProcessCreation0;
    CIT_SPAN_STAT_ITEM  Foreground0;
    CIT_SPAN_STAT_ITEM  Foreground1;
    CIT_SPAN_STAT_ITEM  Foreground2;
    CIT_SPAN_STAT_ITEM  ProcessSuspended;
    CIT_SPAN_STAT_ITEM  ProcessCreation1;
} CIT_USE_DATA_SPAN_STATS;

typedef struct _CIT_SYSTEM_DATA_STATS {
    WORD    Unknown_BootIdRelated0;
    WORD    Unknown_BootIdRelated1;
    WORD    Unknown_BootIdRelated2;
    WORD    Unknown_BootIdRelated3;
    WORD    Unknown_BootIdRelated4;
    WORD    SessionConnects;
    WORD    ProcessForegroundChanges;
    WORD    ContextFlushes;
    WORD    MissingProgData;
    WORD    DesktopSwitches;
    WORD    WinlogonMessage;
    WORD    WinlogonLockHotkey;
    WORD    WinlogonLock;
    WORD    SessionDisconnects;
} CIT_SYSTEM_DATA_STATS;

typedef struct _CIT_USE_DATA_STATS {
    WORD    Crashes;
    WORD    ThreadGhostingChanges;
    WORD    Input;
    WORD    InputKeyboard;
    WORD    Unknown;
    WORD    InputTouch;
    WORD    InputHid;
    WORD    InputMouse;
    WORD    MouseLeftButton;
    WORD    MouseRightButton;
    WORD    MouseMiddleButton;
    WORD    MouseWheel;
} CIT_USE_DATA_STATS;

// PUU
typedef struct _CIT_POST_UPDATE_USE_INFO {
    DWORD   UpdateKey;
    WORD    UpdateCount;
    WORD    CrashCount;
    WORD    SessionCount;
    WORD    LogCount;
    DWORD   UserActiveDurationInS;
    DWORD   UserOrDispActiveDurationInS;
    DWORD   DesktopActiveDurationInS;
    WORD    Version;
    WORD    _Unk0;
    WORD    BootIdMin;
    WORD    BootIdMax;
    DWORD   PMUUKey;
    DWORD   SessionDurationInS;
    DWORD   SessionUptimeInS;
    DWORD   UserInputInS;
    DWORD   MouseInputInS;
    DWORD   KeyboardInputInS;
    DWORD   TouchInputInS;
    DWORD   PrecisionTouchpadInputInS;
    DWORD   InForegroundInS;
    DWORD   ForegroundSwitchCount;
    DWORD   UserActiveTransitionCount;
    DWORD   _Unk1;
    FILETIME    LogTimeStart;
    QWORD   CumulativeUserActiveDurationInS;
    WORD    UpdateCountAccumulationStarted;
    WORD    _Unk2;
    DWORD   BuildUserActiveDurationInS;
    DWORD   BuildNumber;
    DWORD   _UnkDeltaUserOrDispActiveDurationInS;
    DWORD   _UnkDeltaTime;
    DWORD   _Unk3;
} CIT_POST_UPDATE_USE_INFO;

// DP
typedef struct _CIT_DP_MEMOIZATION_ENTRY {
    DWORD   Unk0;
    DWORD   Unk1;
    DWORD   Unk2;
} CIT_DP_MEMOIZATION_ENTRY;

typedef struct _CIT_DP_MEMOIZATION_CONTEXT {
    _CIT_DP_MEMOIZATION_ENTRY   Entries[12];
} CIT_DP_MEMOIZATION_CONTEXT;

typedef struct _CIT_DP_DATA {
    WORD    Version;
    WORD    Size;
    WORD    LogCount;
    WORD    CrashCount;
    DWORD   SessionCount;
    DWORD   UpdateKey;
    QWORD   _Unk0;
    FILETIME    _UnkTime;
    FILETIME    LogTimeStart;
    DWORD   ForegroundDurations[11];
    DWORD   _Unk1;
    _CIT_DP_MEMOIZATION_CONTEXT MemoizationContext;
} CIT_DP_DATA;
"""

c_cit = cstruct.cstruct()
c_cit.load(cit_def)


class CIT:
    def __init__(self, buf):
        compressed_fh = io.BytesIO(buf)
        compressed_size, uncompressed_size = struct.unpack("<2I", compressed_fh.read(8))

        self.buf = lznt1_decompress(compressed_fh)

        self.header = c_cit.CIT_HEADER(self.buf)
        if self.header.MajorVersion != 0x0A:
            raise ValueError("Unsupported CIT version")

        digest = crc32(self.buf[0x14:], crc32(self.buf[:0x10]))
        if self.header.Crc32 != digest:
            raise ValueError("Crc32 mismatch")

        system_data_buf = self.data(self.header.SystemDataOffset, self.header.SystemDataSize, 0x18)
        self.system_data = SystemData(self, c_cit.CIT_PERSISTED(system_data_buf))

        base_use_data_buf = self.data(self.header.BaseUseDataOffset, self.header.BaseUseDataSize, 0x18)
        self.base_use_data = BaseUseData(self, c_cit.CIT_PERSISTED(base_use_data_buf))

        entry_data = self.buf[self.header.EntryDataOffset :]
        self.entries = [Entry(self, entry) for entry in c_cit.CIT_ENTRY[self.header.EntryCount](entry_data)]

    def data(self, offset, size, expected_size=None):
        if expected_size and size > expected_size:
            size = expected_size

        data = self.buf[offset : offset + size]

        if expected_size and size < expected_size:
            data.ljust(expected_size, b"\x00")

        return data

    def iter_bitmap(self, bitmap: bytes):
        bit_delta = timedelta(seconds=self.header.BitPeriodInS)
        ts = wintimestamp(self.header.PeriodStartLocal)

        for byte in bitmap:
            if byte == b"\x00":
                ts += 8 * bit_delta
            else:
                for bit in range(8):
                    if byte & (1 << bit):
                        yield ts
                    ts += bit_delta


class Entry:
    def __init__(self, cit, entry):
        self.cit = cit
        self.entry = entry

        program_buf = cit.data(entry.ProgramDataOffset, entry.ProgramDataSize, 0x1C)
        self.program_data = c_cit.CIT_PROGRAM_DATA(program_buf)

        use_data_buf = cit.data(entry.UseDataOffset, entry.UseDataSize, 0x18)
        self.use_data = ProgramUseData(cit, c_cit.CIT_PERSISTED(use_data_buf))

        self.file_path = None
        self.command_line = None

        if self.program_data.FilePathOffset:
            file_path_buf = cit.data(self.program_data.FilePathOffset, self.program_data.FilePathSize * 2)
            self.file_path = file_path_buf.decode("utf-16-le")

        if self.program_data.CommandLineOffset:
            command_line_buf = cit.data(self.program_data.CommandLineOffset, self.program_data.CommandLineSize * 2)
            self.command_line = command_line_buf.decode("utf-16-le")

    def __repr__(self):
        return f"<Entry file_path={self.file_path!r} command_line={self.command_line!r}>"


class BaseUseData:
    MIN_BITMAPS_SIZE = 0x8
    MIN_SPAN_STATS_SIZE = 0x30
    MIN_STATS_SIZE = 0x18

    def __init__(self, cit, entry):
        self.cit = cit
        self.entry = entry

        bitmap_items = c_cit.CIT_BITMAP_ITEM[entry.BitmapsSize // len(c_cit.CIT_BITMAP_ITEM)](
            cit.data(entry.BitmapsOffset, entry.BitmapsSize, self.MIN_BITMAPS_SIZE)
        )
        bitmaps = [cit.data(item.Offset, item.Size) for item in bitmap_items]
        self.bitmaps = self._parse_bitmaps(bitmaps)
        self.span_stats = self._parse_span_stats(
            cit.data(entry.SpanStatsOffset, entry.SpanStatsSize, self.MIN_SPAN_STATS_SIZE)
        )
        self.stats = self._parse_stats(cit.data(entry.StatsOffset, entry.StatsSize, self.MIN_STATS_SIZE))

    def _parse_bitmaps(self, bitmaps):
        return BaseUseDataBitmaps(self.cit, bitmaps)

    def _parse_span_stats(self, span_stats_data):
        return None

    def _parse_stats(self, stats_data):
        return None


class BaseUseDataBitmaps:
    def __init__(self, cit, bitmaps):
        self.cit = cit
        self._bitmaps = bitmaps

    def _parse_bitmap(self, idx):
        return list(self.cit.iter_bitmap(self._bitmaps[idx]))


class SystemData(BaseUseData):
    MIN_BITMAPS_SIZE = 0x30
    MIN_SPAN_STATS_SIZE = 0x58
    MIN_STATS_SIZE = 0x1C

    def _parse_bitmaps(self, bitmaps):
        return SystemDataBitmaps(self.cit, bitmaps)

    def _parse_span_stats(self, span_stats_data):
        return c_cit.CIT_SYSTEM_DATA_SPAN_STATS(span_stats_data)

    def _parse_stats(self, stats_data):
        return c_cit.CIT_SYSTEM_DATA_STATS(stats_data)


class SystemDataBitmaps(BaseUseDataBitmaps):
    def __init__(self, cit, bitmaps):
        super().__init__(cit, bitmaps)
        self.display_power = self._parse_bitmap(0)
        self.display_request_change = self._parse_bitmap(1)
        self.input = self._parse_bitmap(2)
        self.input_touch = self._parse_bitmap(3)
        self.unknown = self._parse_bitmap(4)
        self.foreground = self._parse_bitmap(5)


class ProgramUseData(BaseUseData):
    def _parse_bitmaps(self, bitmaps):
        return ProgramDataBitmaps(self.cit, bitmaps)

    def _parse_span_stats(self, span_stats_data):
        return c_cit.CIT_USE_DATA_SPAN_STATS(span_stats_data)

    def _parse_stats(self, stats_data):
        return c_cit.CIT_USE_DATA_STATS(stats_data)


class ProgramDataBitmaps(BaseUseDataBitmaps):
    def __init__(self, cit, use_data):
        super().__init__(cit, use_data)
        self.foreground = self._parse_bitmap(0)


# Some inlined utility functions for the purpose of the POC
def wintimestamp(ts, tzinfo=timezone.utc):
    # This is a slower method of calculating Windows timestamps, but works on both Windows and Unix platforms
    # Performance is not an issue for this POC
    return datetime(1970, 1, 1, tzinfo=tzinfo) + timedelta(seconds=float(ts) * 1e-7 - 11644473600)


# LZNT1 derived from https://github.com/google/rekall/blob/master/rekall-core/rekall/plugins/filesystems/lznt1.py
def _get_displacement(offset):
    """Calculate the displacement."""
    result = 0
    while offset >= 0x10:
        offset >>= 1
        result += 1

    return result


DISPLACEMENT_TABLE = array.array("B", [_get_displacement(x) for x in range(8192)])

COMPRESSED_MASK = 1 << 15
SIGNATURE_MASK = 3 << 12
SIZE_MASK = (1 << 12) - 1
TAG_MASKS = [(1 << i) for i in range(0, 8)]


def lznt1_decompress(src):
    """LZNT1 decompress from a file-like object.

    Args:
        src: File-like object to decompress from.

    Returns:
        bytes: The decompressed bytes.
    """
    offset = src.tell()
    src.seek(0, io.SEEK_END)
    size = src.tell() - offset
    src.seek(offset)

    dst = io.BytesIO()

    while src.tell() - offset < size:
        block_offset = src.tell()
        uncompressed_chunk_offset = dst.tell()

        block_header = struct.unpack("<H", src.read(2))[0]
        if block_header & SIGNATURE_MASK != SIGNATURE_MASK:
            break

        hsize = block_header & SIZE_MASK

        block_end = block_offset + hsize + 3

        if block_header & COMPRESSED_MASK:
            while src.tell() < block_end:
                header = ord(src.read(1))
                for mask in TAG_MASKS:
                    if src.tell() >= block_end:
                        break

                    if header & mask:
                        pointer = struct.unpack("<H", src.read(2))[0]
                        displacement = DISPLACEMENT_TABLE[dst.tell() - uncompressed_chunk_offset - 1]

                        symbol_offset = (pointer >> (12 - displacement)) + 1
                        symbol_length = (pointer & (0xFFF >> displacement)) + 3

                        dst.seek(-symbol_offset, io.SEEK_END)
                        data = dst.read(symbol_length)

                        # Pad the data to make it fit.
                        if 0 < len(data) < symbol_length:
                            data = data * (symbol_length // len(data) + 1)
                            data = data[:symbol_length]

                        dst.seek(0, io.SEEK_END)
                        dst.write(data)
                    else:
                        data = src.read(1)
                        dst.write(data)

        else:
            # Block is not compressed
            data = src.read(hsize + 1)
            dst.write(data)

    result = dst.getvalue()
    return result


def print_bitmap(name, bitmap, indent=8):
    print(f"{' ' * indent}{name}:")
    for entry in bitmap:
        print(f"{' ' * (indent + 4)}{entry}")


def print_span_stats(span_stats, indent=8):
    for key, value in span_stats._values.items():
        print(f"{' ' * indent}{key}: {value.Count} times, {value.Duration}ms")


def print_stats(stats, indent=8):
    for key, value in stats._values.items():
        print(f"{' ' * indent}{key}: {value}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("input", type=argparse.FileType("rb"), help="path to SOFTWARE hive file")
    parser.add_argument("--tz", default="UTC", help="timezone to use for parsing local timestamps")
    args = parser.parse_args()

    if not HAS_ZONEINFO:
        print("[!] zoneinfo module not available, falling back to UTC")
        tz = timezone.utc
    else:
        tz = ZoneInfo(args.tz)

    hive = Registry.Registry(args.input)

    try:
        cit_key = hive.open("Microsoft\\Windows NT\\CurrentVersion\\AppCompatFlags\\CIT\\System")
    except Registry.RegistryKeyNotFoundException:
        parser.exit("No CIT\\System key found in the hive specified!")

    for cit_value in cit_key.values():
        data = cit_value.value()

        if len(data) <= 8:
            continue

        print(f"Parsing {cit_value.name()}")
        cit = CIT(data)

        print("Period start:", wintimestamp(cit.header.PeriodStartLocal, tz))
        print("Start time:", wintimestamp(cit.header.StartTimeLocal, tz))
        print("Current time:", wintimestamp(cit.header.CurrentTimeLocal, tz))
        print("Bit period in hours:", cit.header.BitPeriodInS // 60 // 60)
        print("Aggregation period in hours:", cit.header.AggregationPeriodInS // 60 // 60)
        print()

        print("System:")
        print("    Bitmaps:")
        print_bitmap("Display power", cit.system_data.bitmaps.display_power)
        print_bitmap("Display request change", cit.system_data.bitmaps.display_request_change)
        print_bitmap("Input", cit.system_data.bitmaps.input)
        print_bitmap("Input (touch)", cit.system_data.bitmaps.input_touch)
        print_bitmap("Unknown", cit.system_data.bitmaps.unknown)
        print_bitmap("Foreground", cit.system_data.bitmaps.foreground)
        print("    Span stats:")
        print_span_stats(cit.system_data.span_stats)
        print("    Stats:")
        print_stats(cit.system_data.stats)

        print()
        for i, entry in enumerate(cit.entries):
            print(f"Entry {i}:")
            print("    File path:", entry.file_path)
            print("    Command line:", entry.command_line)
            print("    PE TimeDateStamp", datetime.fromtimestamp(entry.program_data.PeTimeDateStamp, tz=timezone.utc))
            print("    PE CheckSum", hex(entry.program_data.PeCheckSum))
            print("    Extra 3:", entry.program_data.Extra3)
            print("    Bitmaps:")
            print_bitmap("Foreground", entry.use_data.bitmaps.foreground)
            print("    Span stats:")
            print_span_stats(entry.use_data.span_stats)
            print("    Stats:")
            print_stats(entry.use_data.stats)
            print()


if __name__ == "__main__":
    main()