# Copyright 2026 SigmaSRC, Inc. All Rights Reserved.

"""Gathering the CPU."""

# To reliably function when invoked interactively as "python ./gather_cpu.py",
# we may have to put ../.. on sys.path.
if __name__ == "__main__":
    import os
    import sys
    try:
        import sigma
    except ImportError:
        sys.path.append(os.path.join(os.pardir, os.pardir))
    del sys
    del os

import re
import os
import logging
from sigma.gather import gathering

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

# Return a list of tuples, where each tuple is a row to add to the
# gather.
def do_windows():
    from win32com.client import GetObject
    slotpattern = re.compile(r"^CPU(?P<slot>[0-9]+)$")
    result = []
    try:
        items = GetObject("WinMgMts://./root/cimv2").InstancesOf("Win32_Processor")
        count = 0
        for item in items:
            cpuid = item.Properties_("DeviceID").Value
            match = slotpattern.match(cpuid)
            if match:
                cpuid = match.group("slot")
            else:
                cpuid = str(count)
            count = count + 1
            result.append((item.Properties_("Name").Value, cpuid))
    except:
        logger.exception("cpu gather failed")
    return result
        
# Take the output from the relevant lshw command as input and
# return a list of rows to add to the gather.  do_linux and do_solaris
# are factored out so we can have a unit test that exercises parsing
# the command output.
def do_linux(lshwout):
    shortcpupat = re.compile(r"^\*-cpu$")
    longcpupat = re.compile(r"^\*-cpu:(?P<slot>[0-9]+).*$")
    productpat = re.compile(r"product: (?P<prod>.*)$")
    result = []
    for line in lshwout.splitlines():
        line = line.lstrip()
        shortmatch = shortcpupat.match(line)
        longmatch = longcpupat.match(line)
        productmatch = productpat.match(line)
        if shortmatch:
            slot = "0"
        elif longmatch:
            slot = longmatch.group("slot")
        elif productmatch:
            result.append((productmatch.group("prod"), slot))
    return result

def do_solaris(psrinfoout):
    slotpat = re.compile(r"processor (?P<slot>[0-9]+) as")
    namepat = re.compile(r"The (?P<name>.*) processor operates")
    result = []
    for line in psrinfoout.splitlines():
        line = line.lstrip()
        slotmatch = slotpat.search(line)
        namematch = namepat.match(line)
        if slotmatch:
            slot = slotmatch.group("slot")
        elif namematch:
            name = namematch.group("name")
            result.append((name, slot))
    return result

def do_osx():
    cpu = ""
    speed = ""
    slots = 0
    regex = re.compile(r"\s*([^:]+):\s*(.*[\S])\s*$")
    pipe = os.popen("/usr/sbin/system_profiler SPHardwareDataType", "r")
    try:
        for line in pipe:
            m = regex.match(line)
            if m:
                key, value = m.groups()
                if key == "CPU Type":
                    cpu = value
                elif key == "CPU Speed":
                    speed = value
                elif key == "Number Of CPUs":
                    slots = int(value)
    finally:
        pipe.close()
    if speed:
        cpu = "%s  %s" % (cpu, speed)
    return [(cpu, i) for i in range(slots)]

def do_aix():
    # XXX prtconf is thorough but slow.
    # Here are two commands that might get us the same info quicker:
    # $ lsdev -c processor       # Gives one line per processor
    # proc0 Available 00-00 Processor
    # $ lsattr -E -l proc0       # Gives details for a processor
    # frequency   1499992000     Processor Speed       False
    # smt_enabled true           Processor SMT enabled False
    # smt_threads 2              Processor SMT threads False
    # state       enable         Processor state       False
    # type        PowerPC_POWER5 Processor type        False
    # $
    regex = re.compile(r"^(\S[^:]*):\s*(\S.*?)\s*$")
    pipe = os.popen("/usr/sbin/prtconf")
    cpu = ""
    speed = ""
    slots = 0
    try:
        for line in pipe:
            m = regex.match(line)
            if m:
                key, value = m.group(1, 2)
                if key == "Processor Type":
                    cpu = value
                elif key == "Processor Clock Speed":
                    speed = value
                elif key == "Number Of Processors":
                    slots = int(value)
    finally:
        pipe.close()
    if speed:
        cpu = "%s  %s" % (cpu, speed)
    return [(cpu, i) for i in range(slots)]

def Gather(collector):
    rows = [] # If all else fails, report nothing
    if os.name == "nt":
        rows = do_windows()
    else:
        ostype = os.uname()[0]
        if ostype == "Linux":
            lshwbin = gathering.get_thirdparty_exe("lshw")
            cmd = lshwbin + " -C processor -disable scsi " + \
                  "-disable ide -disable cpuid"
            rows = do_linux(gathering.cmdoutput(cmd))
        elif ostype == "SunOS":
            rows = do_solaris(gathering.cmdoutput("/usr/sbin/psrinfo -v"))
        elif ostype == "HP-UX":
            rows = []
            # XXX Is the output from model(1) right here?
            p = os.popen("/usr/bin/model")
            try:
                model = p.readline().strip()
            finally:
                p.close()
            p = os.popen("/usr/bin/mpsched -s")
            try:
                for line in p:
                    m = re.match(r"Processor Count\s+:\s+(\d+)", line)
                    if m:
                        count = int(m.group(1))
                        for slot in range(count):
                            rows.append([model, slot])
            finally:
                p.close()
        elif ostype == "Darwin":
            rows = do_osx()
        elif ostype == "AIX":
            rows = do_aix()
        else:
            raise RuntimeError("Don't know how to gather CPU on %s" % ostype)
    table = collector.table("cpu", ["cpu", "slot"])
    for row in rows:
        table.add(row)

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

if __name__ == "__main__":
    _test()
