#
#   Copyright (C) Stonesoft Corporation 2010 - 2012.
#   All rights reserved.
#
#   The StoneGate software, manuals, and technical
#   literature may not be reproduced in any form or
#   by any means except by permission in writing from
#   Stonesoft Corporation.
#

require 'socket'

module P4WebGUI
  @@command_shell_inst = nil
  @@command_shell_mutex = Mutex.new()

  def P4WebGUI::get_command_shell_mutex()
    return @@command_shell_mutex
  end

  def P4WebGUI::get_command_shell_inst()
    return @@command_shell_inst
  end

  def P4WebGUI::set_command_shell_inst( inst )
    @@command_shell_inst = inst
  end

  class P4Exploit
    def initialize( cmdline, log_path, unique_id )
      @cmdline = cmdline
      @keep_executing = true
      @log = []
      @succeeded = false
      @log_path = log_path
      @unique_id = unique_id

      @expect_mutex = Mutex.new()
      @expect_mutex.synchronize do
        @expect = ATFExpect.new( cmdline, 1000 )
      end

      is_predator = false
      if( cmdline =~ /predator4 / or cmdline =~ /evader / )
        is_predator = true;
      end

      logline = "Running exploit with command \"#{cmdline}\""
      @log.push logline
      P4WebGUI::log_web_info( logline )
      
      @watch_thread = Thread.new do
        reply_line = nil
        while @keep_executing
          begin
            reply_line = @expect.expect( /\n/m )
          rescue Errno::EIO
            @expect_mutex.synchronize do
              @expect.pid = nil
              @expect = nil
            end
            break
          rescue Exception => exc
            reply_line = "*timeout*"
          end

          @log.push reply_line.clone()

          # Parse output for some interesting lines
          if reply_line =~ /Info: Listening for command shell control on port (\d+)/
            # Command shell open - replace this line with links to open & close the shell
            cmd_port = $1.to_i()

            # We shall open a shell instance anyway, we need to close the shell gracefully
            P4WebGUI::get_command_shell_mutex().synchronize do
              if not P4WebGUI::get_command_shell_inst().nil?
                puts "A previous command shell instance exists - internal error."
              end

              P4WebGUI::set_command_shell_inst( P4CommandShellControl.new( cmd_port ) )
            end

            @succeeded = true

            open_shell_cmd = "Exploit succeeded! <a class=\"log_link\" href=\"javascript:OpenShell()\">Open shell</a>"+
                             "|"+
                             "<a class=\"log_link\" href=\"javascript:CloseShell()\">Close shell</a>"+
                             "<br>"

            P4WebGUI::log_web( open_shell_cmd, false )
          elsif reply_line =~ /Info: Accepted command shell control connection/
            # No need to log this, internal fluff
          else
            P4WebGUI::log_web( reply_line + "\r\n" )
          end
        end

        rel_log_file_path, rel_record_file_path = write_log()

        if is_predator
          P4WebGUI::log_web( "Finished. Download " +
                           "<a class=\"log_link\" href=\"#{rel_record_file_path}\">traffic capture</a>" +
                           " and " +
                           "<a class=\"log_link\" href=\"#{rel_log_file_path}\">log file</a><hr><br>", false )
        else
          P4WebGUI::log_web( "Finished. Download " +
                           "<a class=\"log_link\" href=\"#{rel_log_file_path}\">log file</a><hr><br>", false )
        end
        
        P4WebGUI::unregister_exploit( self )
        @watch_thread = nil
      end
    end

    # Kill exploit in midflight
    def close()
      puts "P4Exploit::Close() called"
      
      @expect_mutex.synchronize do
        if not @expect.nil?
        	puts "\tKilling @expect"
        	@keep_executing = false
          	@expect.kill()
          	@expect = nil
        end
      end
    end

    # Write the logfile to @unique_id+".log"
    def write_log()

      rel_path = @log_path
      abs_path = "#{$webgui_root_webgui}/public/"

      if @succeeded
        rel_path += "success/"
      else
        rel_path += "failure/"
      end

      rel_path += @unique_id
      abs_path = "#{$webgui_root_webgui}/public/#{rel_path}"

      

      rel_log_file_path = rel_path + ".log"
      abs_log_file_path = abs_path + ".log"
      rel_pcap_file_path = rel_path + ".pcap"
      abs_pcap_file_path = abs_path + ".pcap"

      dirname = File.dirname( abs_pcap_file_path )
      if not File.exists? dirname
        `mkdir -p \"#{dirname}\"`
      end

      # Construct final log filenames
      if File.exists?( abs_log_file_path )
        puts "Log file already found - unique id is not unique after all!"
        P4WebGUI::log_web_info( "Log file already found - unique id is not unique after all!" )
      else
        log_file = File.new( abs_log_file_path, "w" )
        @log.each do |logline|
          log_file.puts logline
        end
        log_file.close()
      end

      # Move temporary capture to correct place
      if File.exists? $webgui_tmp_pcap
        mv_cmd = "mv \"#{$webgui_tmp_pcap}\" \"#{abs_pcap_file_path}\""
        #puts "Running move command #{mv_cmd}"
        mv_output = `#{mv_cmd}`
        #puts "Move output: \"#{mv_output}\""
      end
      
      return rel_log_file_path, rel_pcap_file_path
    end
  end

  # A class that handles communication between the web gui and predator4 tcp shell 
  class P4CommandShellControl
    def initialize( port )
      @shell_port = port
      @socket = nil
      @control_thread = nil

      # Collect payload to send until linefeeds, as shells don't care anyway
      @write_mutex = Mutex.new()
      @write_partial = []

      @read_mutex = Mutex.new()
      @read_payload = []
      @read_payload_offsets = {}

      connect()
    end

    ## Functions called via the web gui

    # Get content as HTML
    def get_content_html( params )
      session_id = params[ "session_id" ]
      if session_id.nil?
        P4WebGUI::log_web_error( "P4CommandShellControl::get_content_html() - session_id not defined!" )
        return ""
      end

      return read( session_id )
    end

    # Write a keypress 
    def write_content( params )
      code = params[ "input" ].to_i()

      if code == 0x0d
        code = [ 0x0d, 0x0a ]
      else
        code = [ code ]
      end

      write( code )
      @read_mutex.synchronize do
        @read_payload.concat code
      end

      return get_content_html( params )
    end

    ## Internal functionality
    def connect()
      if @shell_port < 0 or @shell_port > 0xffff
        raise "Invalid shell port #{@shell_port}"
      end

      begin
        @socket = TCPSocket.new( "127.0.0.1", @shell_port )
       rescue Exception => exc
        return false
      end

      @control_thread = Thread.new do
        payload_str = nil
        while true
          begin
            payload_str = @socket.recv( 0xffff )
          rescue Exception => exc
            # Something broke, exit
            break
          end

          @read_mutex.synchronize do
            @read_payload.concat payload_str.unpack( "C*" )
          end
        end
        @control_thread = nil
      end
    end

    def close()
      if not @socket.nil?
        begin
          @socket.send( "exit\x0d\x0a", 0 )
          @socket.shutdown()
        rescue Exception => exc
          puts "P4CommandShellControl::close() - Got exception #{exc}"
        end
      end

      @control_thread = nil
      
      return true
    end

    def write( keycode_arr )
      @write_mutex.synchronize do
        @write_partial.concat keycode_arr

        if keycode_arr.size() == 2 and keycode_arr == [ 0x0d, 0x0a ]
          begin
            @socket.send( @write_partial.pack( "C*" ), 0 )
          rescue Exception => exc
            puts "P4CommandShellControl::write() - Got exception #{exc}"
          end
          
          @write_partial.clear
        end
      end
    end

    def read( session_id )
      content_str = curr_offset = nil

      @read_mutex.synchronize do
        if @read_payload_offsets[ session_id ].nil?
          @read_payload_offsets[ session_id ] = 0
        end

        curr_offset = @read_payload_offsets[ session_id ]
        if curr_offset > @read_payload.size()
          curr_offset = @read_payload_offsets[ session_id ] = @read_payload.size()
        end

        content_str = @read_payload[ curr_offset..-1 ].pack( "C*" )
        @read_payload_offsets[ session_id ] += content_str.size()
      end
      
      return P4WebGUI::escape_log_web( content_str )
    end
  end
end