# Copyright 2026 SigmaSRC, Inc. All Rights Reserved.

"""Gather script for Windows process info."""


import os
import time
import logging
import tempfile

from sigma.util import console

# All the field names VBS can gather
FIELDNAMES = """
Caption
CommandLine
CreationClassName
CreationDate
CSCreationClassName
CSName
Description
ExecutablePath
ExecutionState
Handle
HandleCount
InstallDate
KernelModeTime
MaximumWorkingSetSize
MinimumWorkingSetSize
Name
OSCreationClassName
OSName
OtherOperationCount
OtherTransferCount
PageFaults
PageFileUsage
ParentProcessId
PeakPageFileUsage
PeakVirtualSize
PeakWorkingSetSize
Priority
PrivatePageCount
ProcessId
QuotaNonPagedPoolUsage
QuotaPagedPoolUsage
QuotaPeakNonPagedPoolUsage
QuotaPeakPagedPoolUsage
ReadOperationCount
ReadTransferCount
SessionId
Status
TerminationDate
ThreadCount
UserModeTime
VirtualSize
WindowsVersion
WorkingSetSize
WriteOperationCount
WriteTransferCount
""".split()

# Subset of field names similar to Unix ps output
SUBSET = """
ProcessId
ParentProcessId
CreationDate
CommandLine
Caption
""".split()

# Template for VBS script; the %s becomes a bunch of Wscript.Echo commands
TEMPLATE = """
strComputer = "."
Set objWMIService = GetObject("winmgmts:\\\\" & strComputer & "\\root\\cimv2")
Set colItems = objWMIService.ExecQuery("Select * from Win32_Process",,48)
For Each objItem in colItems
    %s
    Wscript.Echo "---"
Next
"""

logger = logging.getLogger("gather.gather_winps")

def Gather(collector):
    # write header to output
    table = collector.table("winps", FIELDNAMES)
    for fields in gather_helper(FIELDNAMES):
        table.add(fields)

def Gather_unixps(collector):
    # write header to output
    table = collector.table("ps", "pid ppid stime cmd".split())
    # This is called from gather_ps when os.name == "nt"
    for fields in gather_helper(SUBSET):
        stime = fields[2]
        if stime:
            fields[2] = fixtime(stime)
        if not fields[3]:
            fields[3] = fields[4]
        table.add(fields[:4])

def fixtime(s):
    # Input is of the form YYYYMMDDhhmmss.uuuuuu-zzz
    # Output should be either hh:mm or MonDD
    # Do this in a Python 2.2-on-Windows-compatible way; no strptime():
    assert len(s) == 25
    YYYY, MM, DD, hh, mm, ss = map(int,
                                   [s[0:4], s[4:6], s[6:8],
                                    s[8:10], s[10:12], s[12:14]])
    dot_uuuuuu = s[14:21]
    assert dot_uuuuuu.startswith("."), dot_uuuuuu
    sgn_zzz = s[21:25]
    assert sgn_zzz.startswith("+") or sgn_zzz.startswith("-"), sgn_zzz
    t = (YYYY, MM, DD, hh, mm, ss, -1, -1, -1) # Assume it's local time
    now = time.time()
    if now - 24*3600 < time.mktime(t) <= now:
        return time.strftime("%H:%M", t)
    else:
        return time.strftime("%b%d", t)

def gather_helper(fieldnames):
    if os.name != "nt":
        raise RuntimeError("Can only gather Windows process info on Windows")

    # create the VBS script text
    statements = ["    Wscript.Echo objItem.%s\n" % fn for fn in fieldnames]
    vbscript = TEMPLATE % "".join(statements)

    # create a temp file and stuff the VBS script into it
    tfn = tempfile.mktemp(".vbs")
    try:
        tmp = open(tfn, "w")
        try:
            tmp.write(vbscript)
        finally:
            tmp.close()

        # cscript is the command-line VBS interpreter;
        # //nologo kills off some crufty output
        command = "cscript //nologo " + tfn
        p = console.popen(command, "r")
        rawdata = p.readlines()
        sts = p.close()
        if sts:
            raise RuntimeError("Nonzero exit status from VBS script: %s" % sts)
    finally:
        if os.path.exists(tfn):
            os.remove(tfn)

    # write process entries to output
    ncolumns = len(fieldnames)
    current = []
    for line in rawdata:
        if line == "---\n":
            if len(current) == ncolumns:
                yield current
            else:
                logger.warning("Got the wrong number of columns: %r", current)
            current = []
            continue
        field = line.strip()
        if field == "null":
            field = ""
        current.append(field)
    assert not current, rawdata

def _test():
    from sigma.gather import gathering
    gathering.test(Gather)

if __name__ == "__main__":
    _test()
