lib/noise/pattern.rb in noise-ruby-0.10.0 vs lib/noise/pattern.rb in noise-ruby-0.10.1

@@ -1,25 +1,111 @@
 # frozen_string_literal: true
 
 module Noise
   module Token
-    E = 'e'
-    S = 's'
-    EE = 'ee'
-    ES = 'es'
-    SE = 'se'
-    SS = 'ss'
-    PSK = 'psk'
+    class TokenE
+      def to_s
+        'e'
+      end
+    end
+
+    class TokenS
+      def to_s
+        's'
+      end
+    end
+
+    class TokenDH
+      def mix(symmetric_state, dh_fn, initiator, keypair)
+        private_key, public_key = get_key(keypair, initiator)
+        symmetric_state.mix_key(dh_fn.dh(private_key, public_key))
+      end
+    end
+
+    class TokenEE < TokenDH
+      def get_key(keypair, _initiator)
+        [keypair.e.private_key, keypair.re]
+      end
+
+      def to_s
+        'ee'
+      end
+    end
+
+    class TokenES < TokenDH
+      def get_key(keypair, initiator)
+        initiator ? [keypair.e.private_key, keypair.rs] : [keypair.s.private_key, keypair.re]
+      end
+
+      def to_s
+        'es'
+      end
+    end
+
+    class TokenSE < TokenDH
+      def get_key(keypair, initiator)
+        initiator ? [keypair.s.private_key, keypair.re] : [keypair.e.private_key, keypair.rs]
+      end
+
+      def to_s
+        'se'
+      end
+    end
+    class TokenSS < TokenDH
+      def get_key(keypair, _initiator)
+        [keypair.s.private_key, keypair.rs]
+      end
+
+      def to_s
+        'ss'
+      end
+    end
+    class TokenPSK
+      def to_s
+        'psk'
+      end
+    end
+
+    E = TokenE.new
+    S = TokenS.new
+    EE = TokenEE.new
+    ES = TokenES.new
+    SE = TokenSE.new
+    SS = TokenSS.new
+    PSK = TokenPSK.new
   end
 
+  module Modifier
+    class Psk
+      attr_reader :index
+      def initialize(index)
+        @index = index
+      end
+    end
+
+    class Fallback
+    end
+
+    def self.parse(s)
+      if s.start_with?('psk')
+        index = s.gsub(/psk/, '').to_i
+        Modifier::Psk.new(index)
+      elsif s == 'fallback'
+        Modifier::Fallback.new
+      else
+        raise Noise::Exceptions::UnsupportedModifierError
+      end
+    end
+  end
+
   class Pattern
     attr_reader :tokens, :modifiers, :psk_count, :fallback
 
     def self.create(name)
       pattern_set = name.scan(/([A-Z1]+)([^A-Z]*)/)&.first
       pattern = pattern_set&.first
-      modifiers = pattern_set[1].split('+')
+      modifiers = pattern_set[1].split('+').map { |s| Modifier.parse(s) }
       class_name = "Noise::Pattern#{pattern}"
       klass = Object.const_get(class_name)
       klass.new(modifiers)
     end
 
@@ -29,25 +115,28 @@
       @name = ''
       @psk_count = 0
       @modifiers = modifiers
     end
 
+    def psk?
+      modifiers.any? { |m| m.is_a? Modifier::Psk }
+    end
+
     def apply_pattern_modifiers
       @modifiers.each do |modifier|
-        if modifier.start_with?('psk')
-          index = modifier.gsub(/psk/, '').to_i
+        case modifier
+        when Modifier::Psk
+          index = modifier.index
           raise Noise::Exceptions::PSKValueError if index / 2 > @tokens.size
 
           if index.zero?
             @tokens[0].insert(0, Token::PSK)
           else
             @tokens[index - 1] << Token::PSK
           end
           @psk_count += 1
-        elsif modifier == 'fallback'
+        when Modifier::Fallback
           @fallback = true
-        else
-          raise Noise::Exceptions::UnsupportedModifierError
         end
       end
     end
 
     # initiator [Boolean]