From 635b86d20745bfe3861e61426f3ceaf3a39cf1fc Mon Sep 17 00:00:00 2001
From: AI-Mozi <bieniek.milosz0@gmail.com>
Date: Sun, 1 Sep 2024 13:28:31 +0200
Subject: [PATCH] Add a `Value::WebSocket` class

---
 lib/ronin/recon/values/web_socket.rb | 248 ++++++++++++++++++++
 spec/values/web_socket_spec.rb       | 326 +++++++++++++++++++++++++++
 2 files changed, 574 insertions(+)
 create mode 100644 lib/ronin/recon/values/web_socket.rb
 create mode 100644 spec/values/web_socket_spec.rb

diff --git a/lib/ronin/recon/values/web_socket.rb b/lib/ronin/recon/values/web_socket.rb
new file mode 100644
index 0000000..36a7c41
--- /dev/null
+++ b/lib/ronin/recon/values/web_socket.rb
@@ -0,0 +1,248 @@
+# frozen_string_literal: true
+#
+# ronin-recon - A micro-framework and tool for performing reconnaissance.
+#
+# Copyright (c) 2023-2024 Hal Brodigan (postmodern.mod3@gmail.com)
+#
+# ronin-recon is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# ronin-recon is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with ronin-recon.  If not, see <https://www.gnu.org/licenses/>.
+#
+
+require_relative '../value'
+
+require 'uri'
+
+module Ronin
+  module Recon
+    module Values
+      #
+      # Represents a WebSocket.
+      #
+      # @api public
+      #
+      # @since 0.2.0
+      #
+      class WebSocket < Value
+
+        # The parsed URI.
+        #
+        # @return [URI::HTTP, URI::HTTPS]
+        attr_reader :uri
+
+        #
+        # Initializes the WebSocket value.
+        #
+        # @param [URI::HTTP, URI::HTTPS, String] url
+        #
+        def initialize(url)
+          @uri = URI(url)
+        end
+
+        #
+        # Indicates whether the WebSocket uses `ws://` or `wss://`.
+        #
+        # @return ['ws', 'wss']
+        #
+        def scheme
+          @uri.scheme
+        end
+
+        #
+        # The WebSocket's hostname.
+        #
+        # @return [String]
+        #
+        def host
+          @uri.host
+        end
+
+        #
+        # The WebSocket's port number.
+        #
+        # @return [Integer]
+        #
+        def port
+          @uri.port
+        end
+
+        #
+        # The WebSocket's path.
+        #
+        # @return [String]
+        #
+        def path
+          @uri.path
+        end
+
+        #
+        # The WebSocket's query
+        #
+        # @return [String]
+        #
+        def query
+          @uri.query
+        end
+
+        #
+        # Initializes the 'ws://' WebSocket.
+        #
+        # @param [String] host
+        #   The WebSocket's host.
+        #
+        # @param [Integer] port
+        #   The WebSocket's port.
+        #
+        # @param [String] path
+        #   The WebSocket's path.
+        #
+        # @param [String] query
+        #   The WebSocket's query.
+        #
+        def self.ws(host,port=80,path=nil,query=nil)
+          url = "ws://#{host}:#{port}"
+          url << "#{path}" if path
+          url << "?#{query}" if query
+
+          new(url)
+        end
+
+        #
+        # Initializes the 'wss://' WebSocket.
+        #
+        # @param [String] host
+        #   The WebSocket's host.
+        #
+        # @param [Integer] port
+        #   The WebSocket's port.
+        #
+        # @param [String] path
+        #   The WebSocket's path.
+        #
+        # @param [String] query
+        #   The WebSocket's query.
+        #
+        def self.wss(host,port=443,path=nil,query=nil)
+          url = "wss://#{host}:#{port}"
+          url << "#{path}" if path
+          url << "?#{query}" if query
+
+          new(url)
+        end
+
+        #
+        # Compares the WebSocket to another value.
+        #
+        # @param [Value] other
+        #
+        # @return [Boolean]
+        #
+        def eql?(other)
+          self.class == other.class &&
+            scheme == other.scheme &&
+            host == other.host &&
+            port == other.port &&
+            path == other.path &&
+            query == other.query
+        end
+
+        #
+        # Case equality method used for fuzzy matching.
+        #
+        # @param [Value] other
+        #   The other value to compare.
+        #
+        # @return [Boolean]
+        #   Imdicates whether the other value same as {WebSocket}
+        #
+        def ===(other)
+          case other
+          when WebSocket
+            eql?(other)
+          else
+            false
+          end
+        end
+
+        #
+        # The "hash" value of the WebSocket.
+        #
+        # @return [Integer]
+        #   The hash value of {#scheme}, {#host}, {#port}, {#path} and {#query}.
+        #
+        def hash
+          [self.class, scheme, host, port, path, query].hash
+        end
+
+        # Mapping of {#scheme} values to URI classes.
+        #
+        # @api private
+        URI_CLASSES = {
+          'wss' => URI::WSS,
+          'ws'  => URI::WS
+        }
+
+        #
+        # Converts the WebSocket to URI.
+        #
+        # @return [URI::WS, URI::WSS]
+        #   The URI object for the website.
+        #
+        def to_uri
+          URI_CLASSES.fetch(scheme).build(host: host, port: port, path: path, query: query)
+        end
+
+        #
+        # Converts the WebSocket to a String.
+        #
+        # @return [String]
+        #   The base URL value for the WebSocket.
+        #
+        def to_s
+          @uri.to_s
+        end
+
+        #
+        # Coerces the WebSocket value into JSON.
+        #
+        # @return [Hash{Symbol => Object}]
+        #   The Ruby Hash that will be converted into JSON.
+        #
+        def as_json
+          {
+            type:   :web_socket,
+            scheme: scheme,
+            host:   host,
+            port:   port,
+            path:   path,
+            query:  query
+          }
+        end
+
+        #
+        # Returns the type or kind of recon value.
+        #
+        # @return [:web_socket]
+        #
+        # @note
+        #   This is used internally to map a recon value class to a printable
+        #   type.
+        #
+        # @api private
+        #
+        def self.value_type
+          :web_socket
+        end
+      end
+    end
+  end
+end
diff --git a/spec/values/web_socket_spec.rb b/spec/values/web_socket_spec.rb
new file mode 100644
index 0000000..9a1583b
--- /dev/null
+++ b/spec/values/web_socket_spec.rb
@@ -0,0 +1,326 @@
+require 'spec_helper'
+require 'ronin/recon/values/web_socket'
+
+describe Ronin::Recon::Values::WebSocket do
+  let(:scheme) { 'ws' }
+  let(:host)   { 'example.com' }
+  let(:port)   { 80 }
+  let(:path)   { '/path' }
+  let(:query)  { 'foo=bar' }
+  let(:url)    { "#{scheme}://#{host}:#{port}#{path}?#{query}" }
+
+  subject { described_class.new(url) }
+
+  describe "#initialize" do
+    it "must set #scheme" do
+      expect(subject.scheme).to eq(scheme)
+    end
+
+    it "must set #host" do
+      expect(subject.host).to eq(host)
+    end
+
+    it "must set #port" do
+      expect(subject.port).to eq(port)
+    end
+
+    it "must set #path" do
+      expect(subject.path).to eq(path)
+    end
+
+    it "must set #query" do
+      expect(subject.query).to eq(query)
+    end
+  end
+
+  describe ".wss" do
+    subject { described_class.wss(host,port,path,query) }
+
+    it "must create a WebSocket object with #scheme of 'wss'" do
+      expect(subject.scheme).to eq('wss')
+    end
+
+    it "must set #host" do
+      expect(subject.host).to eq(host)
+    end
+
+    it "must set #port" do
+      expect(subject.port).to eq(port)
+    end
+
+    it "must set #path" do
+      expect(subject.path).to eq(path)
+    end
+
+    it "must set #query" do
+      expect(subject.query).to eq(query)
+    end
+  end
+
+  describe ".ws" do
+    subject { described_class.ws(host,port,path,query) }
+
+    it "must create a WebSocket object with #scheme of 'ws'" do
+      expect(subject.scheme).to eq('ws')
+    end
+
+    it "must set #host" do
+      expect(subject.host).to eq(host)
+    end
+
+    it "must set #port" do
+      expect(subject.port).to eq(port)
+    end
+
+    it "must set #path" do
+      expect(subject.path).to eq(path)
+    end
+
+    it "must set #query" do
+      expect(subject.query).to eq(query)
+    end
+  end
+
+  describe "#eql?" do
+    context "when given an WebSocket object" do
+      context "and the other WebSocket object has the same #scheme, #host, #port, #path and #query" do
+        let(:other) { described_class.new(url) }
+
+        it "must return true" do
+          expect(subject.eql?(other)).to be(true)
+        end
+      end
+
+      context "but the other WebSocket object has a different #scheme" do
+        let(:other_url) { "wss://#{host}:#{port}#{path}?#{query}" }
+        let(:other)     { described_class.new(other_url) }
+
+        it "must return false" do
+          expect(subject.eql?(other)).to be(false)
+        end
+      end
+
+      context "but the other WebSocket object has a different #host" do
+        let(:other_url) { "#{scheme}://other.com:#{port}#{path}?#{query}" }
+        let(:other)     { described_class.new(other_url) }
+
+        it "must return false" do
+          expect(subject.eql?(other)).to be(false)
+        end
+      end
+
+      context "but the other WebSocket object has a different #port" do
+        let(:other_url) { "#{scheme}://#{host}:8000#{path}?#{query}" }
+        let(:other)     { described_class.new(other_url) }
+
+        it "must return false" do
+          expect(subject.eql?(other)).to be(false)
+        end
+      end
+
+      context "but the other WebSocket object has a different #path" do
+        let(:other_url) { "#{scheme}://#{host}:#{port}/not_path?#{query}" }
+        let(:other)     { described_class.new(other_url) }
+
+        it "must return false" do
+          expect(subject.eql?(other)).to be(false)
+        end
+      end
+
+      context "but the other WebSocket object has a different #query" do
+        let(:other_url) { "#{scheme}://#{host}:#{port}#{path}?different=query" }
+        let(:other)     { described_class.new(other_url) }
+
+        it "must return false" do
+          expect(subject.eql?(other)).to be(false)
+        end
+      end
+    end
+
+    context "when given a non-WebSocket object" do
+      let(:other) { Object.new }
+
+      it "must return false" do
+        expect(subject.eql?(other)).to be(false)
+      end
+    end
+  end
+
+  describe "#===" do
+    context "when given an WebSocket object" do
+      context "and the other WebSocket object has the same #scheme, #host, #port, #path and #query" do
+        let(:other) { described_class.new(url) }
+
+        it "must return true" do
+          expect(subject === other).to be(true)
+        end
+      end
+
+      context "but the other WebSocket object has a different #scheme" do
+        let(:other_url) { "wss://#{host}:#{port}#{path}?#{query}" }
+        let(:other)     { described_class.new(other_url) }
+
+        it "must return false" do
+          expect(subject === other).to be(false)
+        end
+      end
+
+      context "but the other WebSocket object has a different #host" do
+        let(:other_url) { "#{scheme}://other.com:#{port}#{path}?#{query}" }
+        let(:other)     { described_class.new(other_url) }
+
+        it "must return false" do
+          expect(subject === other).to be(false)
+        end
+      end
+
+      context "but the other WebSocket object has a different #port" do
+        let(:other_url) { "#{scheme}://#{host}:8000#{path}?#{query}" }
+        let(:other)     { described_class.new(other_url) }
+
+        it "must return false" do
+          expect(subject === other).to be(false)
+        end
+      end
+
+      context "but the other WebSocket object has a different #path" do
+        let(:other_url) { "#{scheme}://#{host}:#{port}/not_path?#{query}" }
+        let(:other)     { described_class.new(other_url) }
+
+        it "must return false" do
+          expect(subject === other).to be(false)
+        end
+      end
+
+      context "but the other WebSocket object has a different #query" do
+        let(:other_url) { "#{scheme}://#{host}:#{port}#{path}?different=query" }
+        let(:other)     { described_class.new(other_url) }
+
+        it "must return false" do
+          expect(subject === other).to be(false)
+        end
+      end
+    end
+
+    context "when given a non-WebSocket object" do
+      let(:other) { Object.new }
+
+      it "must return false" do
+        expect(subject === other).to be(false)
+      end
+    end
+  end
+
+  describe "#hash" do
+    it "must return the #hash of an Array containing the class and the #scheme, #host, and #port" do
+      expect(subject.hash).to eq([described_class, scheme, host, port, path, query].hash)
+    end
+  end
+
+  describe "#to_uri" do
+    context "when the #scheme is wss" do
+      let(:scheme) { 'wss' }
+
+      it "must return URI object with 'wss' scheme" do
+        uri = subject.to_uri
+
+        expect(uri.scheme).to eq('wss')
+        expect(uri.host).to eq(host)
+        expect(uri.port).to eq(port)
+        expect(uri.path).to eq(path)
+        expect(uri.query).to eq(query)
+      end
+    end
+
+    context "when the #scheme is ws" do
+      let(:scheme) { 'ws' }
+
+      it "must return URI object with 'ws' scheme" do
+        uri = subject.to_uri
+
+        expect(uri.scheme).to eq('ws')
+        expect(uri.host).to eq(host)
+        expect(uri.port).to eq(port)
+        expect(uri.path).to eq(path)
+        expect(uri.query).to eq(query)
+      end
+    end
+
+    context "when the #scheme is invalid" do
+      let(:scheme) { 'invalid' }
+
+      it "must raise an error" do
+        expect {
+          subject.to_uri
+        }.to raise_error(KeyError, "key not found: \"invalid\"")
+      end
+    end
+  end
+
+  describe "#to_s" do
+    it "must return a String URL for the WebSocket" do
+      expect(subject.to_s).to eq("#{scheme}://#{host}#{path}?#{query}")
+    end
+
+    context "when the #scheme is 'ws'" do
+      let(:scheme) { 'ws' }
+
+      context "and the #port is 80" do
+        let(:port) { 80 }
+
+        it "must omit the port from the URL String" do
+          expect(subject.to_s).to eq("#{scheme}://#{host}#{path}?#{query}")
+        end
+      end
+
+      context "but the #port is not 80" do
+        let(:port) { 8000 }
+
+        it "must omit the port from the URL String" do
+          expect(subject.to_s).to eq("#{scheme}://#{host}:#{port}#{path}?#{query}")
+        end
+      end
+    end
+
+    context "when the #scheme is 'wss'" do
+      let(:scheme) { 'wss' }
+
+      context "and the #port is 443" do
+        let(:port) { 443 }
+
+        it "must omit the port from the URL String" do
+          expect(subject.to_s).to eq("#{scheme}://#{host}#{path}?#{query}")
+        end
+      end
+
+      context "but the #port is not 443" do
+        let(:port) { 9000 }
+
+        it "must omit the port from the URL String" do
+          expect(subject.to_s).to eq("#{scheme}://#{host}:#{port}#{path}?#{query}")
+        end
+      end
+    end
+  end
+
+  describe "#as_json" do
+    it "must return a Hash containing the type: and scheme:, host:, and port: attributes" do
+      expect(subject.as_json).to eq(
+        {
+          type:   :web_socket,
+          scheme: scheme,
+          host:   host,
+          port:   port,
+          path:   path,
+          query:  query
+        }
+      )
+    end
+  end
+
+  describe ".value_type" do
+    it "must return :web_socket" do
+      expect(described_class.value_type).to eq(:web_socket)
+    end
+  end
+end