vaas_main.rb in vaas-2.0.0

- old
+ new
@@ -12,128 +12,177 @@
 require_relative 'vaas_errors'
 
 module VAAS
   class VaasMain
 
-    attr_accessor :session_id, :websocket, :url
-
-    def initialize(url="wss://gateway-vaas.gdatasecurity.de")
+    def initialize(url = "wss://gateway-vaas.gdatasecurity.de", timeout = 600)
       @url = url
+      @timeout = timeout
+      @requests = {}
     end
 
     def connect(token)
       # connect to endpoint
-      endpoint = Async::HTTP::Endpoint.parse(url, alpn_protocols: Async::HTTP::Protocol::HTTP1.names)
-      self.websocket = Async::WebSocket::Client.connect(endpoint)
+      endpoint = Async::HTTP::Endpoint.parse(@url, alpn_protocols: Async::HTTP::Protocol::HTTP1.names)
+      @websocket = Async::WebSocket::Client.connect(endpoint)
 
-      # send authentication request
-      auth_request = JSON.generate({:kind => "AuthRequest", :token => token})
-      websocket.write(auth_request)
+      read_messages
+      keep_alive
 
-      # receive authentication message
-      while message = websocket.read
-        message = JSON.parse(message)
-        if message['success'] == true
-          self.session_id = message['session_id']
-          break
-        else
-          raise VaasAuthenticationError
+      # send authentication request
+      auth_task = Async do |task|
+        task.with_timeout(@timeout) do
+          @auth_notification = Async::Notification.new
+          auth_request = JSON.generate({:kind => "AuthRequest", :token => token})
+          @websocket.write(auth_request)
+          @websocket.flush
+          @auth_notification.wait
+        rescue Async::TimeoutError
+          close
+          raise VaasTimeoutError
         end
       end
-    end
 
-    def get_authenticated_websocket
-      raise VaasInvalidStateError unless websocket
-      raise VaasInvalidStateError, "connect() was not awaited" unless session_id
-      raise VaasConnectionClosedError if websocket.closed?
-      websocket
+      message = auth_task.wait
+      raise VaasAuthenticationError if message['success'] == false
+      @session_id = message['session_id']
     end
 
-    def close
-        websocket&.close
+    def keep_alive
+      @keep_alive_task = Async do
+        until @websocket.closed?
+          sleep 10
+          @websocket.send_ping
+          @websocket.flush
+        end
+      end
     end
 
-    def __for_sha256(sha256, guid)
-      # send verdict request with sha256
-      websocket = get_authenticated_websocket
-      guid = guid || SecureRandom.uuid.to_s
-      verdict_request =  JSON.generate({:kind => "VerdictRequest",
-                                        :session_id => session_id,
-                                        :sha256 => sha256,
-                                        :guid => guid})
-      websocket.write(verdict_request)
-
-      # receive verdict message
-      while message = websocket.read
-        message = JSON.parse(message)
-        if message['kind'] == "VerdictResponse"
-          return message
+    def read_messages
+      @read_task = Async do |task|
+        task.with_timeout(@timeout) do
+          while message = @websocket.read
+            message = JSON.parse(message)
+            if message['kind'] == "AuthResponse"
+              @auth_notification.signal(message)
+            elsif message['kind'] == "VerdictResponse"
+              @requests[message['guid']].signal(message)
+            end
+          end
+        rescue Async::TimeoutError
+          close
+          raise VaasTimeoutError
         end
       end
     end
 
-    def for_sha256(sha256, guid = nil)
-      response = __for_sha256(sha256, guid)
-      VaasVerdict.new(response)
+    def get_authenticated_websocket
+      raise VaasInvalidStateError unless @websocket
+      raise VaasInvalidStateError, "connect() was not awaited" unless @session_id
+      raise VaasConnectionClosedError if @websocket.closed?
+      @websocket
     end
 
-    def for_file(path, guid = nil)
-      # get sha256 of file and send verdict request
-      sha256 = Digest::SHA256.file(path).hexdigest
-      response = __for_sha256(sha256, guid)
+    def close
+      @keep_alive_task&.stop
+      @read_task&.stop
+      @websocket&.close
+    end
 
-      # upload file if verdict is unknown
-      if response['verdict'] == 'Unknown'
-        upload(response, path)
-      else
-        return VaasVerdict.new(response)
+    def for_sha256(sha256)
+      Async do |task|
+        task.with_timeout(@timeout) do
+          verdict_notification = Async::Notification.new
+          guid = SecureRandom.uuid.to_s
+          websocket = get_authenticated_websocket
+          verdict_request =  JSON.generate({:kind => "VerdictRequest",
+                                            :session_id => @session_id,
+                                            :sha256 => sha256,
+                                            :guid => guid})
+          @requests[guid] = verdict_notification
+          websocket.write(verdict_request)
+          websocket.flush
+          VaasVerdict.new(verdict_notification.wait)
+        rescue Async::TimeoutError
+          close
+          raise VaasTimeoutError
+        end
       end
+    end
 
-      # receive verdict message of uploaded file
-      while message = websocket.read
-        message = JSON.parse(message)
-        if message['kind'] == "VerdictResponse" && message['verdict'] != "Unknown"
-          return VaasVerdict.new(message)
+    def for_url(url)
+      Async do |task|
+        task.with_timeout(@timeout) do
+          verdict_notification = Async::Notification.new
+          guid = SecureRandom.uuid.to_s
+          websocket = get_authenticated_websocket
+          verdict_request =  JSON.generate({:kind => "VerdictRequestForUrl",
+                                            :session_id => @session_id,
+                                            :url => url,
+                                            :guid => guid})
+          @requests[guid] = verdict_notification
+          websocket.write(verdict_request)
+          websocket.flush
+          VaasVerdict.new(verdict_notification.wait)
+        rescue Async::TimeoutError
+          close
+          raise VaasTimeoutError
         end
       end
     end
 
-    def for_url(url, guid = nil)
-      #send verdict request with url
-      websocket = get_authenticated_websocket
-      guid = guid || SecureRandom.uuid.to_s
-      url = URI(url).to_s
-      verdict_request =  JSON.generate({:kind => "VerdictRequestForUrl",
-                                        :session_id => session_id,
-                                        :guid => guid,
-                                        :url => url})
-      websocket.write(verdict_request)
+    def for_file(path)
+      Async do |task|
+        task.with_timeout(@timeout) do
+          sha256 = Digest::SHA256.file(path).hexdigest
+          verdict_notification = Async::Notification.new
+          guid = SecureRandom.uuid.to_s
+          websocket = get_authenticated_websocket
+          verdict_request =  JSON.generate({:kind => "VerdictRequest",
+                                            :session_id => @session_id,
+                                            :sha256 => sha256,
+                                            :guid => guid})
+          @requests[guid] = verdict_notification
+          websocket.write(verdict_request)
+          websocket.flush
+          message = verdict_notification.wait
 
-      # receive verdict message
-      while message = websocket.read
-        message = JSON.parse(message)
-        if message['kind'] == "VerdictResponse"
-          return VaasVerdict.new(message)
+          if message['verdict'] == "Unknown"
+            upload_notification = Async::Notification.new
+            @requests[guid] = upload_notification
+            upload(message, path)
+            VaasVerdict.new(upload_notification.wait)
+          else
+            VaasVerdict.new(message)
+          end
+        rescue Async::TimeoutError
+          close
+          raise VaasTimeoutError
         end
       end
     end
 
-    def upload (message, path)
-      token = message['upload_token']
-      url = message['url']
+    def upload(message, path)
+      Async do |task|
+        task.with_timeout(@timeout) do
+          token = message['upload_token']
+          url = message['url']
 
-      Async do
-        client = Async::HTTP::Internet.new
+          client = Async::HTTP::Internet.new
 
-        header = [['authorization', token]]
-        body = Protocol::HTTP::Body::File.open(File.join(path))
+          header = [['authorization', token]]
+          body = Protocol::HTTP::Body::File.open(File.join(path))
 
-        client.put(url, header, body).read
+          response = client.put(url, header, body)
+          response.read
 
-      rescue => e
-        raise VaasUploadError, e
-      ensure
-        client&.close
+          raise VaasUploadError, "Upload failed with code: #{response.status}" if response.status != 200
+        rescue Async::TimeoutError
+          close
+          raise VaasTimeoutError, "Upload timed out"
+        ensure
+          client&.close
+        end
       end
     end
   end
 end