mtproto 0.0.8 → 0.0.10

data/lib/mtproto/errors.rb

@@ -23,6 +23,14 @@ module MTProto
       @error_message = error_message
       super("RPC Error #{error_code}: #{error_message}")
     end
+
+    def flood_wait?
+      error_message&.start_with?('FLOOD_WAIT_')
+    end
+
+    def flood_wait_seconds
+      error_message&.match(/FLOOD_WAIT_(\d+)/)&.[](1)&.to_i
+    end
   end
 
   class UnexpectedConstructorError < Error

data/lib/mtproto/message/message.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+module MTProto
+  class Message
+    include Binary
+    extend Binary
+
+    attr_reader :auth_key_id, :msg_id, :body,
+      :msg_key, :encrypted_data,
+      :server_salt, :session_id, :seq_no
+
+    def initialize(body, msg_id: self.class.generate_msg_id, auth_key_id: 0)
+      @auth_key_id = auth_key_id
+      @msg_id = msg_id
+      @body = body
+    end
+
+    def self.parse(packet)
+      data = packet.data
+      auth_key_id = b_u64(data[0, 8])
+
+      if auth_key_id.zero?
+        parse_unencrypted(data)
+      else
+        parse_encrypted(data, auth_key_id)
+      end
+    end
+
+    def encrypted?
+      !@auth_key_id.zero?
+    end
+
+    def decrypt(auth_key:, sender: :server)
+      raise 'Message is not encrypted' unless encrypted?
+
+      keys = Crypto::MessageKey.derive_aes_key_iv(auth_key, @msg_key, sender: sender)
+      plaintext = Crypto::AES_IGE.decrypt_ige(@encrypted_data, keys[:aes_key], keys[:aes_iv])
+
+      expected_msg_key = Crypto::MessageKey.generate_msg_key(auth_key, plaintext, sender: sender)
+      raise 'msg_key mismatch!' unless @msg_key == expected_msg_key
+
+      offset = 0
+      @server_salt = plaintext[offset, 8].unpack1('Q<')
+      offset += 8
+      @session_id = plaintext[offset, 8].unpack1('Q<')
+      offset += 8
+      @msg_id = plaintext[offset, 8].unpack1('Q<')
+      offset += 8
+      @seq_no = plaintext[offset, 4].unpack1('L<')
+      offset += 4
+      body_length = plaintext[offset, 4].unpack1('L<')
+      offset += 4
+      @body = plaintext[offset, body_length].bytes
+
+      self
+    end
+
+    def bytes
+      u64_b(@auth_key_id) +
+        u64_b(@msg_id) +
+        u32_b(@body.length) +
+        @body
+    end
+
+    def self.generate_msg_id
+      (Time.now.to_f * (2**32)).to_i & ~3
+    end
+
+    private_class_method def self.parse_unencrypted(data)
+      msg_id = b_u64(data[8, 8])
+      body_length = b_u32(data[16, 4])
+      body = data[20, body_length]
+
+      new(body, msg_id: msg_id, auth_key_id: 0)
+    end
+
+    private_class_method def self.parse_encrypted(data, auth_key_id)
+      msg = allocate
+      msg.instance_variable_set(:@auth_key_id, auth_key_id)
+      msg.instance_variable_set(:@msg_key, data[8, 16].pack('C*'))
+      msg.instance_variable_set(:@encrypted_data, data[24..].pack('C*'))
+      msg
+    end
+  end
+end

data/lib/mtproto/session.rb

@@ -14,7 +14,7 @@ module MTProto
     def next_seq_no(content_related: true)
       current = @seq_no
       @seq_no += content_related ? 1 : 0
-      current * 2 + (content_related ? 1 : 0)
+      (current * 2) + (content_related ? 1 : 0)
     end
   end
 end