From c7afef8cf22aa153970321c341b651a637d4695e Mon Sep 17 00:00:00 2001
From: Bar Hofesh <bar.hofesh@gmail.com>
Date: Tue, 8 Aug 2023 11:55:43 +0300
Subject: [PATCH] v1.6.0

---
 shard.yml                  |  2 +-
 src/sec_tester.cr          |  2 +-
 src/sec_tester/repeater.cr | 86 +++++++++++++++++++++-----------------
 src/sec_tester/scan.cr     |  5 +++
 4 files changed, 54 insertions(+), 41 deletions(-)

diff --git a/shard.yml b/shard.yml
index 01afb42..e326f58 100644
--- a/shard.yml
+++ b/shard.yml
@@ -1,5 +1,5 @@
 name: sec_tester
-version: 1.5.5
+version: 1.6.0
 
 authors:
   - Bar Hofesh <bar.hofesh@gmail.com>
diff --git a/src/sec_tester.cr b/src/sec_tester.cr
index fef20dc..aa6cb25 100644
--- a/src/sec_tester.cr
+++ b/src/sec_tester.cr
@@ -10,7 +10,7 @@ require "./sec_tester/repeater.cr"
 
 module SecTester
   Log     = ::Log.for("SecTester")
-  VERSION = "1.5.5"
+  VERSION = "1.6.0"
 
   backend = ::Log::IOBackend.new(STDOUT)
 
diff --git a/src/sec_tester/repeater.cr b/src/sec_tester/repeater.cr
index 062b229..aa0067e 100644
--- a/src/sec_tester/repeater.cr
+++ b/src/sec_tester/repeater.cr
@@ -8,83 +8,91 @@ module SecTester
     getter id : String = ""
 
     def initialize(@api_key : String, hostname : String = "app.brightsec.com")
-      @socket = SocketIO::Client.new(host: hostname, path: "/api/ws/v1/", namespace: "/workstations")
+      @socket = SocketIO::Client.new(
+        host: hostname,
+        path: "/api/ws/v1/",
+        namespace: "/workstations",
+        decoder: SocketIO::Decoders::MsgpackDecoder.new
+      )
       @socket.connect(
-        data: <<-EOF
-          {"token":"#{@api_key}","domain":"#{System.hostname}"}
-          EOF
+        data: {
+          token:  @api_key,
+          domain: System.hostname,
+        }
       )
       deploy
     end
 
     def deploy
-      @socket.send("\"deploy\"")
+      @socket.emit("deploy")
     end
 
     def close
-      @socket.send("\"undeploy\"")
+      @socket.emit("undeploy")
       @socket.close
     end
 
     def run
-      @socket.on_data do |packet|
-        Log.debug { "Repeater Packet: #{packet}" }
-        data = packet.data
-        case data.as_a[0].as_s
-        when "deployed"
-          @id = data[1]["repeaterId"].as_s
-          Log.debug { "Repeater deployed id: #{@id}" }
-        when "undeployed"
-          @id = ""
-        when "error"
-          Log.error { "Repeater Error: #{data[1]["name"]} - #{data[1]["message"]}" }
-        when "request"
-          if id = packet.id
-            handle_request(data[1], id)
-          end
-        end
+      # handle request events
+      @socket.on("request") do |event|
+        handle_request(event)
+      end
+
+      # hamdle deployed events
+      @socket.on("deployed") do |event|
+        next unless data = event.data
+        @id = data.first["repeaterId"].as_s
+        Log.debug { "Repeater deployed id: #{@id}" }
       end
+
+      # handle undeployed events
+      @socket.on("undeployed") do |event|
+        @id = ""
+      end
+
+      Log.debug { "Repeater started" }
     end
 
-    def handle_request(request : JSON::Any, id : Int64)
-      request_data = request.as_h
+    def handle_request(request_event : SocketIO::Event)
+      # fail fast if no request data or id
+      return unless data = request_event.data
+      return unless request_event.id
+
+      request_data = data[0].try &.as_h?
+      return unless request_data
+
       headers = HTTP::Headers.new
       request_data["headers"].as_h.each do |key, value|
         headers[key] = value.as_s
       end
+      # Make request
       response = HTTP::Client.exec(
         method: request_data["method"].as_s,
         url: request_data["url"].as_s,
         headers: headers,
         body: request_data["body"]?.try &.as_s
       )
+      # prepare response data
       hash = Hash(String, String).new
       response.headers.each do |key, value|
         hash[key] = value.first
       end
+
       data = {
-        "protocol":   "http",
-        "statusCode": response.status_code,
-        "body":       response.body.to_s,
-        "headers":    hash,
+        protocol:   "http",
+        statusCode: response.status_code,
+        body:       response.body.to_s,
+        headers:    hash,
       }
-
-      @socket.send(
-        data: data.to_json,
-        id: id,
-        type: :ack
-      )
+      # send response as ack
+      request_event.ack(data)
     rescue e : Exception
       data = {
         protocol:  "http",
         errorCode: "#{e.class}",
         message:   e.message,
       }
-      @socket.send(
-        data: data.to_json,
-        id: id,
-        type: :ack
-      )
+      request_event.ack(data)
     end
   end
 end
diff --git a/src/sec_tester/scan.cr b/src/sec_tester/scan.cr
index 8118e41..5fcd4b3 100644
--- a/src/sec_tester/scan.cr
+++ b/src/sec_tester/scan.cr
@@ -23,6 +23,11 @@ module SecTester
       spawn do
         @repeater.run
       end
+      while @repeater.id.empty?
+        sleep 1.second
+        Fiber.yield
+        Log.debug { "Waiting for repeater to get ID" }
+      end
     end
 
     private def get_headers : HTTP::Headers