newrelic_security 0.1.0

data/lib/newrelic_security/agent/control/http_context.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+require 'set'
+
+module NewRelic::Security
+  module Agent
+    module Control
+      
+      HTTP_ = 'HTTP_'
+      UNDERSCORE = '_'
+      HYPHEN = '-'
+      REQUEST_METHOD = 'REQUEST_METHOD'
+      PATH_INFO = 'PATH_INFO'
+      RACK_INPUT = 'rack.input'
+      CGI_VARIABLES = ::Set.new(%W[ AUTH_TYPE CONTENT_LENGTH CONTENT_TYPE GATEWAY_INTERFACE HTTPS PATH_INFO PATH_TRANSLATED REQUEST_URI QUERY_STRING REMOTE_ADDR REMOTE_HOST REMOTE_IDENT REMOTE_USER REQUEST_METHOD SCRIPT_NAME SERVER_NAME SERVER_PORT SERVER_PROTOCOL SERVER_SOFTWARE rack.url_scheme ])
+
+      class HTTPContext
+        
+        attr_accessor :time_stamp, :req, :method, :headers, :params, :body, :data_truncated, :route, :cache, :fuzz_files, :event_counter
+
+        def initialize(env)
+          @time_stamp = current_time_millis
+          @req = env.select { |key, _| CGI_VARIABLES.include? key}
+          @method = @req[REQUEST_METHOD]
+          @headers = env.select { |key, _| key.include?(HTTP_) }
+          @headers = @headers.transform_keys{ |key| key[5..-1].gsub(UNDERSCORE, HYPHEN).downcase }
+          request = Rack::Request.new(env) unless env.empty?
+					@params = request&.params
+					@params&.each { |k, v| v.force_encoding(Encoding::UTF_8) if v.is_a?(String) }
+          strio = env[RACK_INPUT]
+          if strio.instance_of?(::StringIO)
+						offset = strio.tell
+						@body = strio.read(NewRelic::Security::Agent.config[:'security.request.body_limit'] * 1024) #after read, offset changes
+						strio.seek(offset)
+            # In case of Grape and Roda strio.read giving empty result, added below approach to handle such cases
+            @body = strio.string if @body.nil? && strio.size > 0
+          elsif defined?(::Rack) && defined?(::Rack::Lint::InputWrapper) && strio.instance_of?(::Rack::Lint::InputWrapper)
+						@body = strio.read(NewRelic::Security::Agent.config[:'security.request.body_limit'] * 1024)
+          elsif defined?(::Protocol::Rack::Input) && defined?(::Protocol::Rack::Input) && strio.instance_of?(::Protocol::Rack::Input)
+            @body = strio.read(NewRelic::Security::Agent.config[:'security.request.body_limit'] * 1024)
+          elsif defined?(::PhusionPassenger::Utils::TeeInput) && strio.instance_of?(::PhusionPassenger::Utils::TeeInput)
+						@body = strio.read(NewRelic::Security::Agent.config[:'security.request.body_limit'] * 1024)
+          end
+          @data_truncated = @body && @body.size >= NewRelic::Security::Agent.config[:'security.request.body_limit'] * 1024
+					strio&.rewind
+					@body = @body.force_encoding(Encoding::UTF_8) if @body.is_a?(String)
+          @cache = Hash.new
+          @fuzz_files = ::Set.new
+          @event_counter = 0
+          NewRelic::Security::Agent.agent.http_request_count.increment
+          NewRelic::Security::Agent.agent.iast_client.completed_requests[@headers[NR_CSEC_PARENT_ID]] = [] if @headers.key?(NR_CSEC_PARENT_ID)
+        end
+
+        def current_time_millis
+          (Time.now.to_f * 1000).to_i
+        end
+
+        def self.get_context
+          ::NewRelic::Agent::Tracer.current_transaction.instance_variable_get(:@security_context_data) if ::NewRelic::Agent::Tracer.current_transaction.instance_variable_defined?(:@security_context_data)
+        end
+
+        def self.set_context(env)
+          ::NewRelic::Agent::Tracer.current_transaction.instance_variable_set(:@security_context_data, HTTPContext.new(env))
+        end
+
+        def self.reset_context
+          ::NewRelic::Agent::Tracer.current_transaction.remove_instance_variable(:@security_context_data) if ::NewRelic::Agent::Tracer.current_transaction.instance_variable_defined?(:@security_context_data)
+        end
+      end
+
+    end
+  end
+end

data/lib/newrelic_security/agent/control/iast_client.rb

@@ -0,0 +1,151 @@
+# frozen_string_literal: true
+require 'net/http'
+require 'json'
+require 'uri'
+require 'set'
+require 'resolv'
+
+module NewRelic::Security
+  module Agent
+    module Control
+      FUZZQ_QUEUE_SIZE = 10000
+      METHOD = 'method'
+      URL = 'url'
+      BODY = 'body'
+      HEADERS = 'headers'
+      VERSION = 'version'
+      IS_GRPC = 'isGrpc'
+      INPUT_CLASS = 'inputClass'
+      SERVER_PORT_1 = 'serverPort'
+      PROBING = 'probing'
+      INTERVAL = 'interval'
+      IS_GRPC_CLIENT_STREAM = 'isGrpcClientStream'
+
+      class IASTClient
+        
+        attr_reader :fuzzQ, :iast_dequeue_threads
+        attr_accessor :cooldown_till_timestamp, :last_fuzz_cc_timestamp, :pending_request_ids, :completed_requests, :iast_data_transfer_request_processor_thread
+
+        def initialize
+          @http = nil
+          @stub = nil
+          @fuzzQ = ::SizedQueue.new(FUZZQ_QUEUE_SIZE)
+          @cooldown_till_timestamp = current_time_millis
+          @last_fuzz_cc_timestamp = current_time_millis
+          @pending_request_ids = ::Set.new
+          @completed_requests = {}
+          create_dequeue_threads
+          create_iast_data_transfer_request_processor
+        end
+  
+        def enqueue(message)
+          @fuzzQ.push(message)
+        rescue ThreadError => error
+          NewRelic::Security::Agent.logger.error "Exception in event enqueue, #{error.inspect}, Dropping fuzz request"
+        end
+
+        private
+
+        def create_dequeue_threads
+          # TODO: Create 3 or more consumers for event sending
+          @iast_dequeue_threads = []
+          3.times do |t|
+            @iast_dequeue_threads << Thread.new do
+              Thread.current.name = "newrelic_security_iast_thread-#{t}"
+              loop do
+                fuzz_request = @fuzzQ.deq #thread blocks when the queue is empty
+                if fuzz_request.request[IS_GRPC]
+                  fire_grpc_request(fuzz_request.id, fuzz_request.request, fuzz_request.reflected_metadata)
+                else
+                  fire_request(fuzz_request.id, fuzz_request.request)
+                end
+                fuzz_request = nil
+              end
+            end
+          end
+        rescue Exception => exception
+          NewRelic::Security::Agent.logger.error "Exception in event queue creation : #{exception.inspect}"
+        end
+        
+        def create_iast_data_transfer_request_processor
+          @iast_data_transfer_request_processor_thread = Thread.new do
+            Thread.current.name = "newrelic_security_iast_data_transfer_request_processor"
+            loop do
+              sleep NewRelic::Security::Agent.config[:policy][VULNERABILITY_SCAN][IAST_SCAN][PROBING][INTERVAL]
+              current_timestamp = current_time_millis
+              cooldown_sleep_time = @cooldown_till_timestamp - current_timestamp
+              sleep cooldown_sleep_time/1000 if cooldown_sleep_time > 0
+              next if current_timestamp - @last_fuzz_cc_timestamp < 5000
+              
+              current_fetch_threshold = 300
+              remaining_record_capacity = @fuzzQ.max
+              current_record_backlog = @fuzzQ.size
+              batch_size = current_fetch_threshold - current_record_backlog
+              if batch_size > 100 && remaining_record_capacity > batch_size
+                iast_data_transfer_request = NewRelic::Security::Agent::Control::IASTDataTransferRequest.new
+                iast_data_transfer_request.batchSize = batch_size * 2
+                iast_data_transfer_request.pendingRequestIds = pending_request_ids.to_a
+                iast_data_transfer_request.completedRequests = completed_requests
+                NewRelic::Security::Agent.agent.event_processor.send_iast_data_transfer_request(iast_data_transfer_request)
+              end
+            end
+          end
+        rescue Exception => exception
+          NewRelic::Security::Agent.logger.error "Exception in create_iast_data_transfer_request_processor creation : #{exception.inspect}"
+        end
+
+        def current_time_millis
+          (Time.now.to_f * 1000).to_i
+        end
+
+        def fire_request(fuzz_request_id, request)
+          unless ::Thread.current[:http]
+            Thread.current[:http] = ::Net::HTTP.new('127.0.0.1', NewRelic::Security::Agent.config[:listen_port])
+            Thread.current[:http].open_timeout = 5
+            if request[PROTOCOL] == HTTPS
+              Thread.current[:http].use_ssl = true
+              Thread.current[:http].verify_mode = OpenSSL::SSL::VERIFY_NONE
+            end
+          end
+          request[HEADERS].delete(VERSION) if request[HEADERS].key?(VERSION)
+          time_before_request = (Time.now.to_f * 1000).to_i
+          response = Thread.current[:http].send_request(request[METHOD], ::URI.parse(request[URL]).to_s, request[BODY], request[HEADERS])
+          time_after_request = (Time.now.to_f * 1000).to_i
+          NewRelic::Security::Agent.logger.debug "IAST fuzz request : time taken : #{time_after_request - time_before_request}ms, #{request.inspect} \nresponse: #{response.inspect}\n"
+        rescue Exception => exception
+          NewRelic::Security::Agent.logger.debug "Unable to fire IAST fuzz request Request : #{request.inspect} Exception : #{exception.inspect} #{exception.backtrace}"
+        ensure
+          NewRelic::Security::Agent.agent.iast_client.completed_requests[fuzz_request_id] = []
+          NewRelic::Security::Agent.agent.iast_client.pending_request_ids.delete(fuzz_request_id)
+        end
+
+        def fire_grpc_request(fuzz_request_id, request, reflected_metadata)
+          service = Object.const_get(request[METHOD].split(SLASH)[0]).superclass
+          method = request[METHOD].split(SLASH)[1]
+          @stub = service.rpc_stub_class.new("localhost:#{request[SERVER_PORT_1]}", :this_channel_is_insecure) unless @stub
+
+          parsed_body =  request[BODY][1..-2].split(',')
+          if reflected_metadata[IS_GRPC_CLIENT_STREAM]
+            chunks_enum = Enumerator.new do |y|
+              parsed_body.each do |b|
+                y << Object.const_get(reflected_metadata[INPUT_CLASS]).decode_json(b)
+              end
+            end
+          else
+            chunks_enum = Object.const_get(reflected_metadata[INPUT_CLASS]).decode_json(request[BODY])
+          end
+          response = @stub.public_send(method, chunks_enum, metadata: request[HEADERS])
+          # response = @stub.send(method, JSON.parse(request['body'], object_class: OpenStruct))
+          # request[HEADERS].delete(VERSION) if request[HEADERS].key?(VERSION)
+          NewRelic::Security::Agent.logger.debug "IAST gRPC client response : #{request.inspect} \n#{response.inspect}\n\n\n\n"
+        rescue Exception => exception
+          NewRelic::Security::Agent.logger.debug "Unable to fire IAST gRPC fuzz request Request : #{request.inspect} Exception : #{exception.inspect} #{exception.backtrace}"
+        ensure
+          NewRelic::Security::Agent.agent.iast_client.pending_request_ids.delete(fuzz_request_id)
+        end
+
+      end
+    end
+  end
+end
+

data/lib/newrelic_security/agent/control/iast_data_transfer_request.rb

@@ -0,0 +1,32 @@
+require 'json'
+require 'set'
+
+module NewRelic::Security
+  module Agent
+    module Control
+      class IASTDataTransferRequest
+        attr_reader :jsonName
+        attr_accessor :batchSize, :pendingRequestIds, :completedRequests
+
+        def initialize
+          @jsonName = :'iast-data-request'
+          @applicationUUID = NewRelic::Security::Agent.config[:uuid]
+          @batchSize = 10
+          @pendingRequestIds = []
+          @completedRequests = Hash.new
+        end
+        
+        def as_json
+          instance_variables.map! do |ivar|
+            [ivar[1..-1].to_sym, instance_variable_get(ivar)]
+          end.to_h
+        end
+
+        def to_json
+          as_json.to_json
+        end
+
+      end
+    end
+  end
+end

data/lib/newrelic_security/agent/control/reflected_xss.rb

@@ -0,0 +1,258 @@
+# frozen_string_literal: true
+require 'set'
+require 'cgi'
+require 'json'
+
+module NewRelic::Security
+  module Agent
+    module Control
+      module ReflectedXSS
+
+        LESS_THAN = '<'
+        GREATER_THAN = '>'
+        EQUAL = '='
+        HTML_COMMENT_START = '!--'
+        HTML_COMMENT_END = '-->'
+        FIVE_COLON = ':::::'
+        SCRIPT = 'script'
+        Content_Type = 'Content-Type'
+        QUERY_STRING = 'QUERY_STRING'
+        REQUEST_URI = 'REQUEST_URI'
+        APPLICATION_JSON = 'application/json'
+        APPLICATION_XML = 'application/xml'
+        APPLICATION_X_WWW_FORM_URLENCODED = 'application/x-www-form-urlencoded'
+        ON1 = 'on'
+        ON2 = 'ON'
+        ON3 = 'On'
+        ON4 = 'oN'
+        SRC ='src'
+        HREF = 'href'
+        ACTION = 'action'
+        FORMACTION = 'formaction'
+        SRCDOC = 'srcdoc'
+        DATA = 'data'
+
+        TAG_NAME_REGEX = ::Regexp.new("<([a-zA-Z_\\-]+[0-9]*|!--)", ::Regexp::MULTILINE | ::Regexp::IGNORECASE )
+        ATTRIBUTE_REGEX = ::Regexp.new("([^(\\/\\s<'\">)]+?)(?:\\s*)=\\s*(('|\")([\\s\\S]*?)(?:(?=(\\\\?))\\5.)*?\\3|.+?(?=\\/>|>|\\?>|\\s|<\\/|$))", Regexp::MULTILINE | Regexp::IGNORECASE)
+        UNSUPPORTED_MEDIA_TYPES = %w[video/ image/ font/ audio/].freeze
+        UNSUPPORTED_CONTENT_TYPES = %w[application/zip application/epub+zip application/gzip application/java-archive application/msword application/octet-stream application/ogg application/pdf application/rtf application/vnd.amazon.ebook application/vnd.apple.installer+xml application/vnd.ms-excel application/vnd.ms-fontobject 
+                                       application/vnd.ms-powerpoint application/vnd.oasis.opendocument.presentation application/vnd.oasis.opendocument.spreadsheet application/vnd.oasis.opendocument.text application/vnd.openxmlformats-officedocument.presentationml.presentation 
+                                       application/vnd.openxmlformats-officedocument.spreadsheetml.sheet application/vnd.openxmlformats-officedocument.wordprocessingml.document application/vnd.rar application/vnd.visio application/x-7z-compressed application/x-abiword application/x-bzip application/x-bzip2 application/x-cdf 
+                                       application/x-freearc application/x-tar application/zip text/calendar ].freeze
+  
+
+        extend self
+
+        def check_xss(http_req, retval)
+          # TODO: Check if enableHTTPRequestPrinting is required.
+          return if http_req.nil? || retval.empty?
+          if retval[1].key?(Content_Type) && (retval[1][Content_Type].start_with?(*UNSUPPORTED_MEDIA_TYPES) || retval[1][Content_Type].start_with?(*UNSUPPORTED_CONTENT_TYPES))
+            return
+          end
+          response_body = ::String.new
+          retval[2].each { |string| response_body << string }
+          construct = check_for_reflected_xss(http_req, retval[1], response_body)
+          NewRelic::Security::Agent.logger.debug "RXSS Attack DATA: #{construct}"
+          if !construct.empty? || NewRelic::Security::Agent::Utils.is_IAST?
+            parameters = Array.new
+            parameters << construct
+            parameters << response_body.force_encoding(ISO_8859_1).encode(UTF_8)
+            NewRelic::Security::Agent::Control::Collector.collect(REFLECTED_XSS, parameters, nil, :response_header => retval[1][Content_Type])
+          end
+        rescue Exception => exception
+          NewRelic::Security::Agent.logger.error "Exception in Reflected XSS detection : #{exception.inspect} #{exception.backtrace}"
+        end
+
+        private
+
+        def check_for_reflected_xss(http_req, headers, response_body)
+          final_attack_construct = ::String.new
+          to_return = ::String.new
+          combined_request_data = decode_request_data(http_req)
+          combined_response_data = decode_response_data(headers, response_body)
+          combined_response_data_string = combined_response_data.to_a.join(FIVE_COLON)
+          attack_constructs  = is_xss(combined_request_data)
+          NewRelic::Security::Agent.logger.debug "RXSS attack_constructs ==> #{attack_constructs}"
+          attack_constructs.each { |construct| to_return = construct if combined_response_data_string.include?(construct) }
+          if !to_return.empty?
+            response_constructs = is_xss(combined_response_data)
+            response_constructs.each { |construct| final_attack_construct = to_return if construct.include?(to_return) }
+          end
+          combined_request_data = nil
+          combined_response_data = nil
+          combined_response_data_string = nil
+          final_attack_construct
+        end
+
+        def decode_request_data(http_req)
+          processed_data = ::Set.new
+          content_type = http_req.req[CONTENT_TYPE]
+          body = http_req.body
+          http_req.req.each do | key, value |
+            process_url_encoded_data_for_xss(processed_data, key)
+            process_url_encoded_data_for_xss(processed_data, value)
+          end          
+          if http_req.params != nil 
+            items = ::Set.new
+            get_key_values(http_req.params, items)
+            items.each { |item| processed_data.add(item) if item.include?(LESS_THAN) }
+          end
+          process_url_encoded_data_for_xss(processed_data, http_req.req[REQUEST_URI])
+          processed_data.add(body) unless body.nil? || body.empty?
+          if body != nil && !body.empty?
+            case content_type
+            when APPLICATION_JSON
+              oldBody = body.dup
+              body = ::JSON.parse(body)
+              if oldBody != body && body.include?(LESS_THAN)
+                processed_data.add(body)
+              end
+            when APPLICATION_XML
+              # Unescaping of xml data is remaining
+              processed_data.add(body)
+            when APPLICATION_X_WWW_FORM_URLENCODED
+              body = ::CGI.unescape(body, UTF_8)
+              processed_data.add(body)
+              oldBody = body
+              body = ::CGI.unescape(body, UTF_8)
+              processed_data.add(body) if oldBody != body && body.include?(LESS_THAN)
+            end
+          end
+          processed_data
+        end
+
+        def decode_response_data(headers, response_body)
+          processed_data = ::Set.new
+          content_type = headers[Content_Type]
+          response_body = response_body
+          processed_body = response_body.force_encoding(UTF_8)
+          processed_data.add(processed_body)
+          old_processed_body = String.new
+          if response_body != nil && !response_body.empty?
+            case content_type
+            when APPLICATION_JSON
+              # do while loop in java code here
+              old_processed_body = processed_body
+              body = ::JSON.parse(processed_body)
+              processed_data.add(body) if old_processed_body != body && body.to_s.include?(LESS_THAN)
+            when APPLICATION_XML
+              # Unescaping of xml data is remaining
+              processed_data.add(processed_data)
+            end
+          end
+          processed_data
+        end
+
+        def is_xss(combined_data)
+          attack_constructs = ::Set.new
+          for data in combined_data do
+            constructs = get_xss_constructs(data)
+            constructs.each { |str| attack_constructs.add(str) }
+          end
+          return attack_constructs
+        end
+
+        def process_url_encoded_data_for_xss(processed_data, data)
+          processed_data.add(data) if data && data.include?(LESS_THAN)
+          decoded_uri = ::CGI.unescape(data, UTF_8) if data
+          processed_data.add(decoded_uri) if decoded_uri && decoded_uri.include?(LESS_THAN)
+        end
+
+        def get_key_values(hash, items)
+          hash.each do |k,v|
+            items.add(k.to_s)
+            if v.instance_of?(Hash)
+              get_key_values(v, items)
+            else
+              items.add(v.to_s) unless v.nil?
+            end
+          end
+        end
+
+        def get_xss_constructs(data)
+          constructs = ::Set.new
+          is_attack_construct = false
+          curr_pos = 0
+          start_pos = 0
+          tmp_curr_pos = 0
+          tmp_start_pos = 0
+      
+          while curr_pos < data.length
+            matcher = TAG_NAME_REGEX.match(data, curr_pos)
+            is_attack_construct = false
+            return constructs if matcher.nil?
+            tagName = matcher[1]
+            return constructs if tagName.empty?
+            start_pos = matcher.begin(0)
+            curr_pos = matcher.end(0) - 1
+            if tagName == HTML_COMMENT_START
+              tmp_curr_pos = start_pos + data.index(HTML_COMMENT_END, start_pos)
+              if tmp_curr_pos == nil
+                break
+              else
+                curr_pos = tmp_curr_pos
+                next
+              end
+            end
+            tmp_start_pos = tmp_curr_pos = data.index(GREATER_THAN, start_pos)
+            tmp_start_pos = start_pos if tmp_curr_pos.nil?
+            while ATTRIBUTE_REGEX.match?(data, curr_pos)
+              attribute_matcher = ATTRIBUTE_REGEX.match(data, curr_pos)
+              attribute_data = attribute_matcher[0]
+              curr_pos = attribute_matcher.end(0) - 1
+              tmp_curr_pos = data.index(GREATER_THAN, tmp_start_pos ? tmp_start_pos : -1)          
+              if tmp_curr_pos == nil || attribute_matcher.begin(0) < tmp_curr_pos
+                tmp_start_pos = tmp_curr_pos = attribute_matcher.end(0) - 1
+                tmp_start_pos += 1
+                if (attribute_matcher[3] == nil || attribute_matcher[3] == EMPTY_STRING) && attribute_matcher.end(0) > tmp_curr_pos
+                  tmp_start_pos = tmp_curr_pos = data.index(GREATER_THAN, attribute_matcher.begin(0)) ? data.index(GREATER_THAN, attribute_matcher.begin(0)) : -1
+                  attribute_data = attribute_data[0..tmp_start_pos]
+                end
+                key = attribute_data[0..attribute_data.index(EQUAL) - 1]
+                val = attribute_data[attribute_data.index(EQUAL) + 1.. attribute_data.length]              
+                if key != nil && key != EMPTY_STRING && key.start_with?(ON1, ON2, ON3, ON4) || key.casecmp?(SRC) || key.casecmp?(HREF) || key.casecmp?(ACTION) || key.casecmp?(FORMACTION) || key.casecmp?(SRCDOC) || key.casecmp?(DATA) || ::CGI.unescapeHTML(val).gsub(/[[:space:]]/, EMPTY_STRING).match?(/javascript:/i)
+                  is_attack_construct = true
+                end
+              else
+                break
+              end
+            end
+            if tmp_curr_pos != nil && tmp_curr_pos > 0
+              curr_pos = tmp_curr_pos
+            end
+            if data[curr_pos] != GREATER_THAN
+              tmp = data.index(GREATER_THAN, curr_pos)
+              if tmp != nil
+                curr_pos = tmp
+              elsif !is_attack_construct
+                next
+              end
+            end
+        
+            if tagName.strip.casecmp?(SCRIPT)
+              location_of_end_tag = data.index(/<\/script/i, curr_pos)
+              if location_of_end_tag != nil
+                body = data[curr_pos + 1..location_of_end_tag-1]
+                if body != nil && body != EMPTY_STRING
+                  constructs.add(data[start_pos..curr_pos] + body)
+                  next
+                end
+              else
+                body = data[curr_pos + 1 ..  data.length]
+                tag_end = body.index(GREATER_THAN)
+                if body != nil && body != EMPTY_STRING && tag_end != nil
+                  body = body[tag_end..data.length]
+                  constructs.add(data[start_pos..curr_pos] + body)
+                  break
+                end
+              end
+            end
+          constructs.add(data[start_pos..curr_pos]) if is_attack_construct
+          end
+          constructs
+        end
+
+      end
+    end
+  end
+end

data/lib/newrelic_security/agent/control/websocket_client.rb

@@ -0,0 +1,131 @@
+# frozen_string_literal: true
+
+require 'newrelic_security/websocket-client-simple/client'
+require 'openssl'
+require 'singleton'
+
+module NewRelic::Security
+  module Agent
+    module Control
+
+      NR_CSEC_CONNECTION_TYPE = 'NR-CSEC-CONNECTION-TYPE'
+      NR_LICENSE_KEY = 'NR-LICENSE-KEY'
+      NR_AGENT_RUN_TOKEN = 'NR-AGENT-RUN-TOKEN'
+      NR_CSEC_VERSION = 'NR-CSEC-VERSION'
+      NR_CSEC_COLLECTOR_TYPE = 'NR-CSEC-COLLECTOR-TYPE'
+      NR_CSEC_BUILD_NUMBER = 'NR-CSEC-BUILD-NUMBER'
+      NR_CSEC_MODE = 'NR-CSEC-MODE'
+      NR_CSEC_APP_UUID = 'NR-CSEC-APP-UUID'
+      NR_CSEC_JSON_VERSION = 'NR-CSEC-JSON-VERSION'
+      NR_ACCOUNT_ID = 'NR-ACCOUNT-ID'
+      NR_CSEC_ENTITY_NAME = 'NR-CSEC-ENTITY-NAME'
+      NR_CSEC_ENTITY_GUID = 'NR-CSEC-ENTITY-GUID'
+      NR_CSEC_IAST_DATA_TRANSFER_MODE = 'NR-CSEC-IAST-DATA-TRANSFER-MODE'
+
+      class WebsocketClient
+        include Singleton
+
+        attr_accessor :ws
+
+        def connect()
+
+          headers = Hash.new
+          headers[NR_CSEC_CONNECTION_TYPE] = LANGUAGE_COLLECTOR
+          headers[NR_LICENSE_KEY] = NewRelic::Security::Agent.config[:license_key]
+          headers[NR_AGENT_RUN_TOKEN] = NewRelic::Security::Agent.config[:agent_run_id]
+          headers[NR_CSEC_VERSION] = NewRelic::Security::VERSION
+          headers[NR_CSEC_COLLECTOR_TYPE] = RUBY
+          headers[NR_CSEC_BUILD_NUMBER] = '0000'
+          headers[NR_CSEC_MODE] = NewRelic::Security::Agent.config[:mode]
+          headers[NR_CSEC_APP_UUID] = NewRelic::Security::Agent.config[:uuid]
+          headers[NR_CSEC_JSON_VERSION] = NewRelic::Security::Agent.config[:json_version]
+          headers[NR_ACCOUNT_ID] = NewRelic::Security::Agent.config[:account_id]
+          headers[NR_CSEC_ENTITY_NAME] = NewRelic::Security::Agent.config[:app_name]
+          headers[NR_CSEC_ENTITY_GUID] = NewRelic::Security::Agent.config[:entity_guid]
+          headers[NR_CSEC_IAST_DATA_TRANSFER_MODE] = PULL
+
+          begin
+            cert_store = ::OpenSSL::X509::Store.new
+            cert_store.add_cert ::OpenSSL::X509::Certificate.new(::IO.read("#{__dir__}/../resources/cert.pem"))
+            NewRelic::Security::Agent.logger.info "Websocket connection URL : #{NewRelic::Security::Agent.config[:validator_service_url]}"
+            connection = NewRelic::Security::WebSocket::Client::Simple.connect NewRelic::Security::Agent.config[:validator_service_url], headers: headers, cert_store: cert_store
+            @ws = connection
+
+            connection.on :open do
+              NewRelic::Security::Agent.logger.debug "Websocket connected with IC, AgentEventMachine #{NewRelic::Security::Agent::Utils.filtered_log(connection.inspect)}"
+              NewRelic::Security::Agent.init_logger.info "[STEP-4] => Web socket connection to SaaS validator established successfully"
+              NewRelic::Security::Agent.agent.event_processor.send_app_info
+              NewRelic::Security::Agent.agent.event_processor.send_application_url_mappings
+              NewRelic::Security::Agent.config.enable_security
+            end
+        
+            connection.on :message do |msg|
+              if msg.type == :ping
+                connection.send(EMPTY_STRING, :type => :pong)
+              elsif msg.type == :text
+                # NewRelic::Security::Agent.logger.debug "Received IC Agent Message: #{msg.data.inspect}"
+                ControlCommand.handle_ic_command(msg.data)
+              end
+            end
+
+            connection.on :close do |e|
+              NewRelic::Security::Agent.logger.info "Closing websocket connection : #{e.inspect}\n"
+              NewRelic::Security::Agent.config.disable_security
+              Thread.new { NewRelic::Security::Agent.agent.reconnect(0) } if e
+            end
+
+            connection.on :error do |e|
+              NewRelic::Security::Agent.logger.error "Error in websocket connection : #{e.inspect} #{e.backtrace}"
+              Thread.new { NewRelic::Security::Agent::Control::WebsocketClient.instance.close(true) }
+            end
+          rescue Errno::EPIPE => exception
+            NewRelic::Security::Agent.logger.error "Unable to connect to validator_service: #{exception.inspect}"
+            NewRelic::Security::Agent.config.disable_security
+          rescue Errno::ECONNRESET => exception
+            NewRelic::Security::Agent.logger.error "Unable to connect to validator_service: #{exception.inspect}"
+            NewRelic::Security::Agent.config.disable_security
+            Thread.new { NewRelic::Security::Agent.agent.reconnect(15) }
+          rescue Errno::ECONNREFUSED => exception
+            NewRelic::Security::Agent.logger.error "Unable to connect to validator_service: #{exception.inspect}"
+            NewRelic::Security::Agent.config.disable_security
+            Thread.new { NewRelic::Security::Agent.agent.reconnect(15) }
+          rescue => exception
+            NewRelic::Security::Agent.logger.error "Exception in websocket init: #{exception.inspect} #{exception.backtrace}"
+            NewRelic::Security::Agent.config.disable_security
+            Thread.new { NewRelic::Security::Agent.agent.reconnect(15) }
+          end
+          headers = nil
+        end
+
+        def send(message)
+          message_json = message.to_json
+          NewRelic::Security::Agent.logger.debug "Sending #{message.jsonName} : #{message_json}"
+          res = @ws.send(message_json)
+          if res && message.jsonName == :Event
+            NewRelic::Security::Agent.agent.event_sent_count.increment
+            if NewRelic::Security::Agent::Utils.is_IAST_request?(message.httpRequest[:headers])
+              NewRelic::Security::Agent.agent.iast_event_stats.sent.increment
+            else
+              NewRelic::Security::Agent.agent.rasp_event_stats.sent.increment
+            end
+          end
+          NewRelic::Security::Agent.agent.exit_event_stats.sent.increment if res && message.jsonName == :'exit-event'
+        rescue Exception => exception
+          NewRelic::Security::Agent.logger.error "Exception in sending message : #{exception.inspect} #{exception.backtrace}"
+          NewRelic::Security::Agent.agent.event_drop_count.increment if message.jsonName == :Event
+          NewRelic::Security::Agent.agent.event_processor.send_critical_message(exception.message, "SEVERE", caller_locations[0].to_s, Thread.current.name, exception)
+        end
+
+        def close(reconnect = true)
+          @ws.close(reconnect) if @ws
+        end
+
+        def is_open?
+          return @ws.open? if @ws
+          false
+        end
+
+      end
+    end
+  end
+end