srum_query.py

from __future__ import print_function

import argparse

import ctypes
from datetime import datetime
import collections
import json
import os.path
import csv

import windows
import windows.generated_def as gdef
from windows.generated_def import EnumValue, EnumType, FlagMapper
from windows.utils import  datetime_from_systemtime

## == The live extraction of SRUM data ==
# https://velociraptor.velocidex.com/digging-into-the-system-resource-usage-monitor-srum-afbadb1a375
# https://docs.velociraptor.app/blog/2019/2019-12-31_digging-into-the-system-resource-usage-monitor-srum-afbadb1a375/
# https://www.forensafe.com/blogs/srum.html
# Looks like some columns may contains dates & co.
# Do we want to extract these as well ?
# https://github.com/libyal/esedb-kb/blob/main/documentation/System%20Resource%20Usage%20Monitor%20(SRUM).asciidoc

## SRU Definition (generated by PythonForWindows)

SruAppIdExePath = EnumValue("_SRU_APPID_TYPE", "SruAppIdExePath", 0x0)
SruAppIdSvcName = EnumValue("_SRU_APPID_TYPE", "SruAppIdSvcName", 0x1)
SruAppIdPkgId = EnumValue("_SRU_APPID_TYPE", "SruAppIdPkgId", 0x2)


class _SRU_APPID_TYPE(EnumType):
    values = [SruAppIdExePath, SruAppIdSvcName, SruAppIdPkgId]
    mapper = FlagMapper(*values)

SRU_APPID_TYPE = _SRU_APPID_TYPE

SruColTypeULong = EnumValue("_SRU_STATS_COLUMN_TYPE", "SruColTypeULong", 0x0)
SruColTypeULong64 = EnumValue("_SRU_STATS_COLUMN_TYPE", "SruColTypeULong64", 0x1)
SruColTypeString = EnumValue("_SRU_STATS_COLUMN_TYPE", "SruColTypeString", 0x2)
SruColTypeBinary = EnumValue("_SRU_STATS_COLUMN_TYPE", "SruColTypeBinary", 0x3)

class _SRU_STATS_COLUMN_TYPE(EnumType):
    values = [SruColTypeULong, SruColTypeULong64, SruColTypeString, SruColTypeBinary]
    mapper = FlagMapper(*values)


SRU_STATS_COLUMN_TYPE = _SRU_STATS_COLUMN_TYPE

class _SRU_APPID(ctypes.Structure):
    _fields_ = [
        ("AppIdType", SRU_APPID_TYPE),
        ("Name", gdef.LPWSTR),
    ]

SRU_APPID = _SRU_APPID

class _SRU_SID(ctypes.Structure):
    _fields_ = [
        ("Size", gdef.PVOID),
        ("Sid", gdef.PSID),
    ]


SRU_SID = _SRU_SID

class _SRU_BINARY_DATA(ctypes.Structure):
    _fields_ = [
        ("DataLength", gdef.UINT),
        ("Binary", gdef.PBYTE),
    ]

class _SRU_STATS_COLUMN_DATA(ctypes.Union):
    _fields_ = [
        ("Ulong", gdef.ULONG),
        ("Ulong64", gdef.ULONG64),
        ("String", gdef.LPWSTR),
        ("Binary", _SRU_BINARY_DATA),
    ]


SRU_STATS_COLUMN_DATA = _SRU_STATS_COLUMN_DATA


class _SRU_STATS_COLUMN(ctypes.Structure):
    _fields_ = [
        ("ColId", gdef.USHORT),
        ("ColType", SRU_STATS_COLUMN_TYPE),
        ("ColData", SRU_STATS_COLUMN_DATA),
    ]


PSRU_STATS_COLUMN = ctypes.POINTER(_SRU_STATS_COLUMN)
SRU_STATS_COLUMN = _SRU_STATS_COLUMN


class _SRU_STATS_RECORD(ctypes.Structure):
    _fields_ = [
        ("TimeStamp", gdef.SYSTEMTIME),
        ("AppId", SRU_APPID),
        ("UserId", SRU_SID),
        ("ColCount", gdef.SHORT),
        ("Columns", PSRU_STATS_COLUMN),
    ]


SRU_STATS_RECORD = _SRU_STATS_RECORD
PSRU_STATS_RECORD = ctypes.POINTER(_SRU_STATS_RECORD)


class _SRU_STATS_RECORD_SET(ctypes.Structure):
    _fields_ = [
        ("RecordCount", gdef.USHORT),
        ("RecordList", PSRU_STATS_RECORD),
    ]


SRU_STATS_RECORD_SET = _SRU_STATS_RECORD_SET
PSRU_STATS_RECORD_SET = ctypes.POINTER(_SRU_STATS_RECORD_SET)


SRUM_PROVIDER_PER_NAME = {
                            'SruProviderClassNetwork': 0,
                            'SruProviderClassCpu': 1,
                            'SruProviderClassNotification': 2,
                            'SruProviderClassClassicApp': 3,
                            'SruProviderClassEnergy': 4,
                            'SruProviderClassNetworkConnectivity': 5,
                            'SruProviderClassBatterySense': 6,
                            'SruProviderClassEnergyEstimation': 7,
                            'SruProviderClassTaggedEnergy': 8,
                            'SruProviderClassAppTimeline': 9,
                            'SruProviderClassVelocityFeature': 10,
                            'SruProviderClassSdpCpu': 11,
                            'SruProviderClassSdpEventLog': 12,
                            'SruProviderClassSdpNetwork': 13,
                            'SruProviderClassSdpPerfCounter': 14,
                            'SruProviderClassSdpPhysicalDisk': 15,
                            'SruProviderClassSdpVolume': 16
                        }

# Flags for SruQueryStatsEx
# looks like both of these can yield some different results

SruQueryStatsFlagsAllUsers = 1
SruQueryStatsGlobalOnly = 2

now = datetime.now()

### SRUM Collector code
SRUM_RECORD_FIELD = [
    "appname",
    "apptype",
    "date",
    "sid",
    "columns"
]

SrumRecord = collections.namedtuple("SrumRecord", SRUM_RECORD_FIELD)
SrumColumn = collections.namedtuple("SrumColumn", ["index", "type", "data"])

all_error = set()


def is_srum_dll_available():
    try:
        dll = ctypes.windll.srumapi
    except WindowsError:
        return False
    return True

SRUM_DB_FILE = r"C:\Windows\System32\sru\SRUDB.dat"

def is_srum_db_present():
    return os.path.exists(SRUM_DB_FILE)


## Standard call

COLTYPE_TO_FIEDLNAME = {
    SruColTypeULong: "Ulong",
    SruColTypeULong64: "Ulong64",
    SruColTypeString: "String",
    SruColTypeBinary: "Binary",
}

def normalize_api_record(record):
    columns = []
    for rcol in record.Columns[:record.ColCount]:
        coltype = rcol.ColType
        targetfield = COLTYPE_TO_FIEDLNAME.get(coltype.value, "ERROR")
        coldata = getattr(rcol.ColData, targetfield, "ERROR")
        if coltype.value == SruColTypeBinary:
            coldata = tuple(coldata.Binary[:coldata.DataLength])
        columns.append(SrumColumn(index=rcol.ColId, type=coltype.value, data=coldata))

    return SrumRecord(date=datetime_from_systemtime(record.TimeStamp),
                      appname=record.AppId.Name,
                      apptype=record.AppId.AppIdType.value,
                      sid=str(record.UserId.Sid) if record.UserId.Size else None,
                      columns=tuple(columns)
                      )

def collect_srum_data_from_provider_via_api(providername, flags, start_date, end_date):
    providerid = SRUM_PROVIDER_PER_NAME[providername]
    provider_result = []
    for collect_flag in flags:
        psru_record_set = PSRU_STATS_RECORD_SET()
        print("[~] SruQueryStatsEx prov={} flag={}".format(providername, collect_flag))
        ret = ctypes.windll.srumapi.SruQueryStatsEx(providerid,
                                                ctypes.byref(start_date),
                                                ctypes.byref(end_date),
                                                collect_flag,
                                                None,
                                                ctypes.byref(psru_record_set))
        # 0 -> ERROR_SUCCESS
        # Operation may succed but leave sru_record_set at None
        if (ret != gdef.ERROR_SUCCESS) or not psru_record_set:
            print("  Empty (return value: {})".format(ret))
            continue

        try: # Try Finally to release the sru_record_set pointer in any case
            record_set = psru_record_set[0]
            print("  RecordCount: {}".format(record_set.RecordCount))
            for i in range(record_set.RecordCount):
                srumdata = normalize_api_record(record_set.RecordList[i])
                provider_result.append(srumdata)
        finally:
            ctypes.windll.srumapi.SruFreeRecordSet(psru_record_set)
    return provider_result


def systemtime_from_datetime(dtime):
    return gdef.SYSTEMTIME(wYear=dtime.year, wMonth=dtime.month, wDay=dtime.day, wHour=dtime.hour, wMinute=dtime.minute, wSecond=dtime.second, wMilliseconds=dtime.microsecond // 1000)


def collect_srum_data(args):
    db_present = is_srum_db_present()
    dll_present = is_srum_dll_available()
    print("[~] is_srum_db_present: {0}".format(db_present))
    print("[~] is_srum_dll_available: {0}".format(dll_present))
    if not dll_present:
        raise OSError(" [-] DLL <srumapi> not present, cannot query")


    srum_start = systemtime_from_datetime(args.start_date)
    srum_end = systemtime_from_datetime(args.end_date)

    srum_data = {}
    for providername in args.providers:
        srum_data[providername] = collect_srum_data_from_provider_via_api(providername, args.flags, srum_start, srum_end)
    return srum_data


def display_stats_for_srum_data(srumdata):
    total_first_date = None
    total_last_date= None
    total_count = 0
    print("======== STATISTICS ========")
    for provider, proventries in srumdata.items():
        first_date = None
        last_date = None
        count = 0
        if not proventries:
            print("Provider: {0} : EMPTY".format(provider))
            continue
        for entry in proventries:
            count += 1
            if first_date is None or entry.date < first_date:
                first_date = entry.date
            if last_date is None or entry.date > last_date:
                last_date = entry.date
        # Update total
        total_count += count
        if first_date is not None and (total_first_date is None or first_date < total_first_date):
            total_first_date = first_date
        if last_date is not None and (total_last_date is None or last_date > total_last_date):
            total_last_date = last_date
        print("--- {} ---".format(provider))
        print("* Nb entries: {}".format(count))
        print("* First date: {}".format(first_date))
        print("* Last date: {}".format(last_date))

    print("--- Global ---")
    print("* Nb entries: {}".format(total_count))
    print("* First date: {}".format(total_first_date))
    print("* Last date: {}".format(total_last_date))
    return

def filter_srum_data(srumdata, filter):
    for provider, proventries in srumdata.items():
        if not proventries:
            continue
        srumdata[provider] = [e for e in proventries if e.appname is not None and filter in e.appname]
    return srumdata

def write_outfile_json(filename, srumdata):
    dictsrumdata = {}
    # Transform srumdata to dict for better json results
    for provider, entries in srumdata.items():
        dictsrumdata[provider] = dictentries = []
        for entry in entries:
            dictentry = entry._asdict()
            dictentry["date"] = dictentry["date"].isoformat()
            dictentry["columns"] = [c._asdict() for c in dictentry["columns"]]
            dictentries.append(dictentry)

    # Dump JSON
    with open(filename, "w", encoding="utf-8") as fout:
        json.dump(dictsrumdata, fout, indent=4)
    print("[+] Json written in <{}>".format(filename))

def write_outfile_csv(filename, srumdata):
    fieldnames = ["provider", "appname", "apptype", "date", "sid"]
    with open(filename, 'w', newline='', encoding="utf-8") as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        for provider, entries in srumdata.items():
            for entry in entries:
                writer.writerow({'provider': provider,
                                 "appname": entry.appname,
                                 "apptype": entry.apptype.name,
                                 "date": entry.date.isoformat(),
                                 "sid": entry.sid,
                                 })

    print("[+] CSV written in <{}>".format(filename))

def main():
    parser = argparse.ArgumentParser(prog=__file__)
    parser.add_argument('--providers', nargs="+", choices=SRUM_PROVIDER_PER_NAME, metavar="PROVIDER_NAME", help="The list of providers to query", default="all")
    parser.add_argument('--flags', nargs=1, choices=["GLOBAL", "ALLUSERS", "BOTH"], help="The flags to try", default="BOTH")
    parser.add_argument('--start_date', default="1970-01-01")
    parser.add_argument('--end_date', default="now")
    parser.add_argument('--filter', help="Keep only entries that contain FILTER (Python-side filtering)")
    parser.add_argument('--outfile', default="srum.json", help="Output file (Allowed extensions are .json & .csv) default is 'srum.json'")

    args = parser.parse_args()
    if args.end_date == "now":
        args.end_date = datetime.now()

    args.start_date = datetime.fromisoformat(args.start_date)

    if args.flags == "BOTH":
        args.flags = [SruQueryStatsFlagsAllUsers, SruQueryStatsGlobalOnly]
    elif args.flags == "GLOBAL":
        args.flags = [SruQueryStatsGlobalOnly]
    else: # allusers
        args.flags = [SruQueryStatsFlagsAllUsers]

    if args.providers == "all":
        args.providers = list(SRUM_PROVIDER_PER_NAME)

    windows.com.init() # CoInitialize
    srumdata = collect_srum_data(args)
    if args.filter:
        srumdata = filter_srum_data(srumdata, args.filter)

    if args.outfile.endswith(".json"):
        write_outfile_json(args.outfile, srumdata)
    elif args.outfile.endswith(".csv"):
        write_outfile_csv(args.outfile, srumdata)
    else:
        raise ValueError("Unknown extension for outfile {} (known are '.json' & '.csv')".format(args.outfile))


    display_stats_for_srum_data(srumdata)


if __name__ == "__main__":
    main()