# Copyright 2026 SigmaSRC, Inc. All Rights Reserved.

"""Gather script for media info."""

import os
import platform
import logging
import re
if os.name == "nt":
    from win32com.client import GetObject
    import winreg as _winreg

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

basicos = platform.uname()[0]
if basicos in [ "Microsoft", "Microsoft Windows" ]:
    basicos = "Windows"

if basicos != "Windows":
    raise Exception("Can't get media info on this operating system: %r" %
                    basicos)

# this fn gets a list of CD GUID and CD drive type as a tuple for current user
def GetCDDriveType():
    driveList = []
    keyname = "Software\\Microsoft\\Windows\\CurrentVersion\\Explorer\\CD Burning\\Drives"
    try:
        root = _winreg.ConnectRegistry(None, _winreg.HKEY_USERS)
        try:
            nsubkeys, nvalues, mtime = _winreg.QueryInfoKey(root)
            for i in range(nsubkeys):
                try:
                    userkeyname = _winreg.EnumKey(root, i)
                except OSError as err:
                    continue 
                # ignore 'local system', 'local service' and 'network service' accounts
                if userkeyname.startswith("S-1-5-18") or userkeyname.startswith("S-1-5-19") or userkeyname.startswith("S-1-5-20"): 
                    continue
                # ignore classes for any user and the default user subkey
                if userkeyname.lower().endswith("_classes") or userkeyname.lower() == ".default": continue
                guidskeyname = "%s\\%s" % (userkeyname, keyname)
                try:
                    guidskey = _winreg.OpenKey(root, guidskeyname)
                except OSError as err:
                    continue 
                try:
                    msubkeys, mvalues, mmtime = _winreg.QueryInfoKey(guidskey)
                    for j in range(msubkeys):
                        akeyname = _winreg.EnumKey(guidskey, j)
                        aguidkeyname = "%s\\%s\\%s" % (userkeyname, keyname, akeyname)
                        #guidList.append(aguidkeyname)
                        try:
                            finalkey = _winreg.OpenKey(root, aguidkeyname)
                        except OSError as err:
                            continue
                        try:
                            xsubkeys, xvalues, xmtime = _winreg.QueryInfoKey(finalkey)
                            for i in range(xvalues):
                                nm, val, typ = _winreg.EnumValue(finalkey, i)
                                if not isinstance(nm, str):
                                    nm = str(nm)
                                if not isinstance(val, str):
                                    val = str(val)
                                if nm.lower() == "drive type":
                                    driveList.append((akeyname, val))
                                    break
                        finally:
                            _winreg.CloseKey(finalkey)
                finally:
                    _winreg.CloseKey(guidskey)
        finally:
            _winreg.CloseKey(root)
    except:
        pass
    return driveList


# this fn gets a list of CD GUID and CD drive-name as a tuple for current user
def GetCDDriveName():
    driveList = []
    keyname = "Software\\Microsoft\\Windows\\CurrentVersion\\Explorer\\MountPoints2\\CPC\\Volume"
    try:
        root = _winreg.ConnectRegistry(None, _winreg.HKEY_USERS)
        try:
            nsubkeys, nvalues, mtime = _winreg.QueryInfoKey(root)
            for i in range(nsubkeys):
                try:
                    userkeyname = _winreg.EnumKey(root, i)
                except OSError as err:
                    continue 
                # ignore 'local system', 'local service' and 'network service' accounts
                if userkeyname.startswith("S-1-5-18") or userkeyname.startswith("S-1-5-19") or userkeyname.startswith("S-1-5-20"): 
                    continue
                # ignore classes for any user and the default user subkey
                if userkeyname.lower().endswith("_classes") or userkeyname.lower() == ".default": continue
                guidskeyname = "%s\\%s" % (userkeyname, keyname)
                try:
                    guidskey = _winreg.OpenKey(root, guidskeyname)
                except OSError as err:
                    continue 
                try:
                    msubkeys, mvalues, mmtime = _winreg.QueryInfoKey(guidskey)
                    for j in range(msubkeys):
                        akeyname = _winreg.EnumKey(guidskey, j)
                        aguidkeyname = "%s\\%s\\%s" % (userkeyname, keyname, akeyname)
                        #guidList.append(aguidkeyname)
                        try:
                            finalkey = _winreg.OpenKey(root, aguidkeyname)
                        except OSError as err:
                            continue
                        try:
                            xsubkeys, xvalues, xmtime = _winreg.QueryInfoKey(finalkey)
                            for i in range(xvalues):
                                nm, val, typ = _winreg.EnumValue(finalkey, i)
                                u = val.decode("utf-16")
                                if nm.lower() == "data":                                 
                                    s = ""
                                    # get the printable characters from the id
                                    for i in range(0, len(u)):
                                        if ord(u[i]) < 128 and ord(u[i]) > 31:
                                            s = s + u[i]
                                    driveList.append((akeyname, s.lower()))
                                    break
                        finally:
                            _winreg.CloseKey(finalkey)
                finally:
                    _winreg.CloseKey(guidskey)
        finally:
            _winreg.CloseKey(root)
    except:
        pass
    return driveList


# checks if the CD is writable
def CheckWritableCD(medname, meddesc):
    # get the list of all CD drives with CD Name and type as a tuple
    driveList = []
    driveType = GetCDDriveType()
    driveName = GetCDDriveName()
    for aType in driveType:
        for aName in driveName:
            if aType[0].find(aName[0]) >= 0:
                driveList.append((aName[1], aType[1]))

    pnpDeviceID = "dummy"
    desc = meddesc
    writable = 'no'
    try:
        objItems = GetObject("WinMgMts://./root/cimv2").InstancesOf("Win32_CDROMDrive")
        for item in objItems:
            drivename = item.Properties_("Drive").Value
            if drivename != medname:
                continue
            desc = item.Properties_("Name").Value
            pnpDeviceID = item.Properties_("PNPDeviceID").Value
            s1 = str(pnpDeviceID)
            s2 = s1.replace("\\", "#").lower()
            pnpDeviceID = s2.replace("/", "#")
            break
    except:
        logger.exception("media gather: getting writable cd list failed")

    # check if CD is writable
    if driveList == []:
        regexp = r"(CD+ROM|CD-ROM|CDROM|CD ROM)"
        if re.search(regexp, desc):
            writable = 'no'
        else:
            regexp = r"(CD R|CD RW|CD RAM|CD+R|CD+R/RW|CD+RW|CD-R|CD-R/RW|CD-RW|CDRW|CD-RAM|CD+RAM|CDRAM)"
            if re.search(regexp, desc):
                writable = 'yes'
            else:
                writable = 'no'
    else:
        for aDrive in driveList:
            s = aDrive[0]
            if s.find(pnpDeviceID) >= 0:
                if aDrive[1] == 1 or aDrive[1] == 2:
                    writable = 'yes'
    return writable


# checks if a DVD is writable
def CheckWritableDVD(desc):
    regexp = r"(DVD+ROM|DVD-ROM|DVDROM|DVD ROM)"
    if re.search(regexp, desc):
        dvdWritable = 'no'
    else:
        regexp = r"(DVD R|DVD RW| DVD RAM|DVD+R|DVD+R/RW|DVD+RW|DVD-R|DVD-R/RW|DVD-RW|DVDRW|DVD-RAM|DVD+RAM|DVDRAM)"
        if re.search(regexp, desc):
            dvdWritable = 'yes'
        else:
            dvdWritable = 'no'
    return dvdWritable

    
def Gather(collector):
    table = collector.table("media", ["medtype", "medname", "description", "writable", "totalsize", "usedspace"]) 
    # enumerate all disks                
    try:
        objItems = GetObject("WinMgMts://./root/cimv2").InstancesOf("Win32_LogicalDisk")
        for item in objItems:
            # initialize all var for this item
            writable = 'yes'
            addDVD = False
            strtype = "Unknown"

            # get all the required properties
            medname = item.Properties_("Name").Value
            desc = item.Properties_("Description").Value
            size = item.Properties_("Size").Value
            freespace = item.Properties_("FreeSpace").Value
            medtype = item.Properties_("DriveType").Value

            # size and freespace may be 'None' when CD is not loaded
            if size is None:
                size = 0
            else:
                size = int(size)
            if freespace is None:
                freespace = 0
            else:
                freespace = int(freespace)

            # get the user friendly names and the writable property
            # for each media type
            if medtype == 1:
                strtype = "No Root Directory"
                writable = 'unknown'
            elif medtype == 2:
                strtype = "Removable Disk"
            elif medtype == 3:
                strtype = "Local Disk"
            elif medtype == 4:
                strtype = "Network Drive"
                writable = 'unknown'
            elif medtype == 6:
                strtype = "RAM Disk" 
            elif medtype == 5:
                strtype = "CD"
                # check if CD is writable
                writable = CheckWritableCD(medname, desc)
                # check for DVD and if it's writable
                regexp = r"(\sDVD|/DVD|\\DVD)"
                if re.search(regexp, desc):
                    addDVD = True
                    dvdWritable = CheckWritableDVD(desc)

            # add a line for this media            
            table.add((strtype, medname, desc, writable, size, size - freespace))
            if addDVD:
                table.add(("DVD", medname, desc, dvdWritable, size, size - freespace))
    except:
        logger.exception("media gather: getting drive list failed")


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

if __name__ == "__main__":
    _test()
