Add support for AES decryption.

Co-authored-by: Robert Haines <rhaines@manchester.ac.uk>
Reviewed-by: @Ph0tonic

Author: jplot

.rubocop_todo.yml

@@ -13,6 +13,7 @@ Gemspec/DevelopmentDependencies:
 Lint/MissingSuper:
   Exclude:
     - 'lib/zip/extra_field.rb'
+    - 'lib/zip/extra_field/aes.rb'
     - 'lib/zip/extra_field/ntfs.rb'
     - 'lib/zip/extra_field/old_unix.rb'
     - 'lib/zip/extra_field/universal_time.rb'
@@ -23,7 +24,7 @@ Lint/MissingSuper:
 # Offense count: 5
 # Configuration parameters: CountComments, CountAsOne.
 Metrics/ClassLength:
-  Max: 650
+  Max: 675
 
 # Offense count: 21
 # Configuration parameters: IgnoredMethods.
@@ -60,6 +61,7 @@ Naming/AccessorMethodName:
 Style/ClassAndModuleChildren:
   Exclude:
     - 'lib/zip/extra_field/generic.rb'
+    - 'lib/zip/extra_field/aes.rb'
     - 'lib/zip/extra_field/ntfs.rb'
     - 'lib/zip/extra_field/old_unix.rb'
     - 'lib/zip/extra_field/universal_time.rb'

lib/zip.rb

@@ -30,6 +30,7 @@
 require 'zip/crypto/encryption'
 require 'zip/crypto/null_encryption'
 require 'zip/crypto/traditional_encryption'
+require 'zip/crypto/aes_encryption'
 require 'zip/inflater'
 require 'zip/deflater'
 require 'zip/streamable_stream'

lib/zip/crypto/aes_encryption.rb

@@ -0,0 +1,116 @@
+# frozen_string_literal: true
+
+require 'openssl'
+
+module Zip
+  module AESEncryption # :nodoc:
+    VERIFIER_LENGTH = 2
+    BLOCK_SIZE = 16
+    AUTHENTICATION_CODE_LENGTH = 10
+
+    VERSION_AE_1 = 0x01
+    VERSION_AE_2 = 0x02
+
+    VERSIONS = [
+      VERSION_AE_1,
+      VERSION_AE_2
+    ].freeze
+
+    STRENGTH_128_BIT = 0x01
+    STRENGTH_192_BIT = 0x02
+    STRENGTH_256_BIT = 0x03
+
+    STRENGTHS = [
+      STRENGTH_128_BIT,
+      STRENGTH_192_BIT,
+      STRENGTH_256_BIT
+    ].freeze
+
+    BITS = {
+      STRENGTH_128_BIT => 128,
+      STRENGTH_192_BIT => 192,
+      STRENGTH_256_BIT => 256
+    }.freeze
+
+    KEY_LENGTHS = {
+      STRENGTH_128_BIT => 16,
+      STRENGTH_192_BIT => 24,
+      STRENGTH_256_BIT => 32
+    }.freeze
+
+    SALT_LENGTHS = {
+      STRENGTH_128_BIT => 8,
+      STRENGTH_192_BIT => 12,
+      STRENGTH_256_BIT => 16
+    }.freeze
+
+    def initialize(password, strength)
+      @password = password
+      @strength = strength
+      @bits = BITS[@strength]
+      @key_length = KEY_LENGTHS[@strength]
+      @salt_length = SALT_LENGTHS[@strength]
+    end
+
+    def header_bytesize
+      @salt_length + VERIFIER_LENGTH
+    end
+
+    def gp_flags
+      0x0001
+    end
+  end
+
+  class AESDecrypter < Decrypter # :nodoc:
+    include AESEncryption
+
+    def decrypt(encrypted_data)
+      @hmac.update(encrypted_data)
+
+      idx = 0
+      decrypted_data = +''
+      amount_to_read = encrypted_data.size
+
+      while amount_to_read.positive?
+        @cipher.iv = [@counter + 1].pack('Vx12')
+        begin_index = BLOCK_SIZE * idx
+        end_index = begin_index + [BLOCK_SIZE, amount_to_read].min
+        decrypted_data << @cipher.update(encrypted_data[begin_index...end_index])
+        amount_to_read -= BLOCK_SIZE
+        @counter += 1
+        idx += 1
+      end
+
+      decrypted_data
+    end
+
+    def reset!(header)
+      raise Error, "Unsupported encryption AES-#{@bits}" unless STRENGTHS.include? @strength
+
+      salt = header[0...@salt_length]
+      pwd_verify = header[-VERIFIER_LENGTH..]
+      key_material = OpenSSL::KDF.pbkdf2_hmac(
+        @password,
+        salt:       salt,
+        iterations: 1000,
+        length:     (2 * @key_length) + VERIFIER_LENGTH,
+        hash:       'sha1'
+      )
+      enc_key = key_material[0...@key_length]
+      enc_hmac_key = key_material[@key_length...(2 * @key_length)]
+      enc_pwd_verify = key_material[-VERIFIER_LENGTH..]
+
+      raise Error, 'Bad password' if enc_pwd_verify != pwd_verify
+
+      @counter = 0
+      @cipher = OpenSSL::Cipher::AES.new(@bits, :CTR)
+      @cipher.decrypt
+      @cipher.key = enc_key
+      @hmac = OpenSSL::HMAC.new(enc_hmac_key, OpenSSL::Digest.new('SHA1'))
+    end
+
+    def check_integrity(auth_code)
+      raise Error, 'Integrity fault' if @hmac.digest[0...AUTHENTICATION_CODE_LENGTH] != auth_code
+    end
+  end
+end

lib/zip/crypto/decrypted_io.rb

@@ -4,9 +4,11 @@ module Zip
   class DecryptedIo # :nodoc:all
     CHUNK_SIZE = 32_768
 
-    def initialize(io, decrypter)
+    def initialize(io, decrypter, compressed_size)
       @io = io
       @decrypter = decrypter
+      @offset = io.tell
+      @compressed_size = compressed_size
     end
 
     def read(length = nil, outbuf = +'')
@@ -18,6 +20,10 @@ def read(length = nil, outbuf = +'')
         buffer << produce_input
       end
 
+      if @decrypter.kind_of?(::Zip::AESDecrypter) && input_finished?
+        @decrypter.check_integrity(@io.read(::Zip::AESEncryption::AUTHENTICATION_CODE_LENGTH))
+      end
+
       outbuf.replace(buffer.slice!(0...(length || buffer.bytesize)))
     end
 
@@ -31,12 +37,17 @@ def buffer
       @buffer ||= +''
     end
 
+    def pos
+      @io.tell - @offset
+    end
+
     def input_finished?
-      @io.eof
+      @io.eof || pos >= @compressed_size
     end
 
     def produce_input
-      @decrypter.decrypt(@io.read(CHUNK_SIZE))
+      chunk_size = [CHUNK_SIZE, @compressed_size - pos].min
+      @decrypter.decrypt(@io.read(chunk_size))
     end
   end
 end

lib/zip/entry.rb

@@ -203,6 +203,10 @@ def zip64?
       !@extra['Zip64'].nil?
     end
 
+    def aes?
+      !@extra['AES'].nil?
+    end
+
     def file_type_is?(type) # :nodoc:
       ftype == type
     end
@@ -382,6 +386,7 @@ def read_local_entry(io) # :nodoc:
 
       read_extra_field(extra, local: true)
       parse_zip64_extra(true)
+      parse_aes_extra
       @local_header_size = calculate_local_header_size
     end
 
@@ -511,6 +516,7 @@ def read_c_dir_entry(io) # :nodoc:
       check_c_dir_entry_comment_size
       set_ftype_from_c_dir_entry
       parse_zip64_extra(false)
+      parse_aes_extra
     end
 
     def file_stat(path) # :nodoc:
@@ -800,6 +806,20 @@ def parse_zip64_extra(for_local_header) # :nodoc:
       end
     end
 
+    def parse_aes_extra # :nodoc:
+      return unless aes?
+
+      if @extra['AES'].vendor_id != 'AE'
+        raise Error, "Unsupported encryption method #{@extra['AES'].vendor_id}"
+      end
+
+      unless ::Zip::AESEncryption::VERSIONS.include? @extra['AES'].vendor_version
+        raise Error, "Unsupported encryption style #{@extra['AES'].vendor_version}"
+      end
+
+      @compression_method = @extra['AES'].compression_method if ftype != :directory
+    end
+
     # For DEFLATED compression *only*: set the general purpose flags 1 and 2 to
     # indicate compression level. This seems to be mainly cosmetic but they are
     # generally set by other tools - including in docx files. It is these flags

lib/zip/extra_field.rb

@@ -91,6 +91,7 @@ def local_size
 require 'zip/extra_field/unix'
 require 'zip/extra_field/zip64'
 require 'zip/extra_field/ntfs'
+require 'zip/extra_field/aes'
 
 # Copyright (C) 2002, 2003 Thomas Sondergaard
 # rubyzip is free software; you can redistribute it and/or

lib/zip/extra_field/aes.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+module Zip
+  # Info-ZIP Extra for AES encryption
+  class ExtraField::AES < ExtraField::Generic # :nodoc:
+    attr_reader :vendor_version, :vendor_id, :encryption_strength, :compression_method
+
+    HEADER_ID = [0x9901].pack('v')
+    register_map
+
+    def initialize(binstr = nil)
+      @vendor_version = nil
+      @vendor_id = nil
+      @encryption_strength = nil
+      @compression_method = nil
+      binstr && merge(binstr)
+    end
+
+    def ==(other)
+      @vendor_version == other.vendor_version &&
+        @vendor_id == other.vendor_id &&
+        @encryption_strength == other.encryption_strength &&
+        @compression_method == other.compression_method
+    end
+
+    def merge(binstr)
+      return if binstr.empty?
+
+      size, content = initial_parse(binstr)
+      (size && content) || return
+
+      @vendor_version, @vendor_id,
+        @encryption_strength, @compression_method = content.unpack('va2Cv')
+    end
+
+    def pack_for_local
+      [@vendor_version, @vendor_id,
+       @encryption_strength, @compression_method].pack('va2Cv')
+    end
+
+    def pack_for_c_dir
+      pack_for_local
+    end
+  end
+end

lib/zip/input_stream.rb

@@ -154,7 +154,20 @@ def get_decrypted_io # :nodoc:
       header = @archive_io.read(@decrypter.header_bytesize)
       @decrypter.reset!(header)
 
-      ::Zip::DecryptedIo.new(@archive_io, @decrypter)
+      compressed_size =
+        if @current_entry.incomplete? && @current_entry.crc == 0 &&
+           @current_entry.compressed_size == 0 && @complete_entry
+          @complete_entry.compressed_size
+        else
+          @current_entry.compressed_size
+        end
+
+      if @decrypter.kind_of?(::Zip::AESDecrypter)
+        compressed_size -= @decrypter.header_bytesize
+        compressed_size -= ::Zip::AESEncryption::AUTHENTICATION_CODE_LENGTH
+      end
+
+      ::Zip::DecryptedIo.new(@archive_io, @decrypter, compressed_size)
     end
 
     def get_decompressor # :nodoc: