# frozen_string_literal: true

require 'openssl'

# Extends Ruby's standard OpenSSL module with the CCM (Counter with CBC-MAC) class.
#
# This module is part of Ruby's standard library and is only reopened here
# to provide support for the CCM authenticated encryption mode (as defined in RFC 3610).
module OpenSSL
  # CCMError used for wrong parameter resonse.
  class CCMError < StandardError
  end

  # Abstract from http://tools.ietf.org/html/rfc3610:
  #
  # Counter with CBC-MAC (CCM) is a generic authenticated encryption
  # block cipher mode.  CCM is defined for use with 128-bit block
  # ciphers, such as the Advanced Encryption Standard (AES).
  #
  # At the moment there is no update function, because length of
  # data and additional_data are needed at the begin of cipher process.
  # In future init(nonce, data_len, additional_data_len) could
  # be a solution, to solve this problem. After init, update(data)
  # could be used to set additional_data first followed by data.
  class CCM
    # Searches for supported algorithms within OpenSSL
    #
    # @return [[String]] supported algorithms
    def self.ciphers
      selected = OpenSSL::Cipher.ciphers.select { |c| c.match(/-(128|192|256)-CBC$/i) }
      @ciphers ||= selected.map { |e| e[0..-9].upcase }.uniq
    end

    # Creates a new CCM object.
    #
    # @param cipher [String] one of the supported algorithms like 'AES'
    # @param key [String] the key used for encryption and decryption
    # @param mac_len [Number] the length of the mac.
    #        needs to be in 4, 6, 8, 10, 12, 14, 16
    #
    # @return [Object] the new CCM object
    def initialize(cipher, key, mac_len)
      raise CCMError, "unsupported cipher algorithm (#{cipher})" unless CCM.ciphers.include?(cipher.upcase)
      raise CCMError, 'invalid key length' unless key.b.length >= 16
      raise CCMError, 'invalid mac length' unless (4..16).step(2).include?(mac_len)

      cipher_key_size = if key.length < 24 then '128'
                        elsif key.length < 32 then '192'
                        else '256' # rubocop:disable Lint/ElseLayout
                        end

      @cipher = OpenSSL::Cipher.new("#{cipher.upcase}-#{cipher_key_size}-CBC")
      @key = key
      @mac_len = mac_len
    end

    # Encrypts the input data and appends mac for authentication.
    # If there is additional data, its included into mac calculation.
    #
    # @param data [String] the data to encrypt
    # @param nonce [String] the nonce used for encryption
    # @param additional_data [String] additional data to
    #        authenticate with mac (not part of the output)
    #
    # @return [String] the encrypted data with appended mac
    def encrypt(data, nonce, additional_data = '')
      valid?(data, nonce, additional_data)

      crypt(data, nonce) + mac(data, nonce, additional_data)
    end

    # Decrypts the input data and checks the appended mac.
    # If additional data was used for encryption, its needed
    # for decryption, to check the authentication (mac).
    #
    # @param data [String] the data to decrypt
    # @param nonce [String] the nonce used for decryption
    # @param additional_data [String] additional data to check
    #        authentication (not part of the output)
    #
    # @return [String] the decrypted data without mac
    def decrypt(data, nonce, additional_data = '')
      valid?(data, nonce, additional_data)

      new_data = crypt(data.b[0...-@mac_len], nonce)
      new_mac = mac(new_data, nonce, additional_data)
      return new_data if new_mac == data.b[-@mac_len..-1]

      ''
    end

    private

    def valid?(data, nonce, additional_data)
      raise CCMError, 'invalid nonce length' unless (7..13).include?(nonce.b.length)
      raise CCMError, 'invalid data length' unless data.b.length < 2**(8 * (15 - nonce.b.length))
      raise CCMError, 'invalid additional_data length' unless additional_data.b.length < 2**64

      true
    end

    def crypt(data, nonce)
      result = String.new
      data.bytes.each_slice(16).with_index(1) do |block, b|
        counter = get_counter(nonce, b).bytes
        block.length.times { |i| counter[i] ^= block[i] }
        result << counter[0, block.length].pack('C*')
      end
      result
    end

    def mac(data, nonce, additional_data)
      @cipher.reset
      @cipher.encrypt
      @cipher.key = @key
      @cipher.iv = "\x00" * 16

      mac = init_with_b0(data, nonce, additional_data)
      !additional_data.empty? && mac = process(additional_data)
      !data.empty? && mac = @cipher.update(data + padding(data)).bytes[-16..-1]

      a0 = get_counter(nonce, 0).bytes
      16.times { |i| mac[i] ^= a0[i] }
      mac[0...@mac_len].pack('C*')
    end

    def init_with_b0(data, nonce, additional_data)
      b0 = Array.new(8, 0)
      b0[0] = (additional_data.empty? ? 0 : 64) \
             + (8 * ((@mac_len - 2) / 2)) \
             + (14 - nonce.b.length)
      b0 += [data.b.length].pack('Q').reverse.bytes
      b0[1, nonce.b.length] = nonce.bytes
      @cipher.update(b0.pack('C*')).bytes
    end

    def process(additional_data)
      len = additional_data.b.length
      d = case
          when len < (2**16) - (2**8)
            [len].pack('n')
          when len < 2**32
            "\xFF\xFE#{[len].pack('N')}"
          else
            "\xFF\xFF#{[len].pack('Q').reverse}"
          end + additional_data
      @cipher.update(d + padding(d)).bytes[-16..-1]
    end

    def padding(data)
      return '' if (data.b.length % 16).zero?

      "\x00" * (16 - (data.b.length % 16))
    end

    def get_counter(nonce, index)
      a = Array.new(8, 0)
      a[0] = 14 - nonce.b.length
      a += [index].pack('Q').reverse.bytes
      a[1, nonce.b.length] = nonce.bytes

      @cipher.reset
      @cipher.encrypt
      @cipher.key = @key
      @cipher.iv = "\x00" * 16
      @cipher.update(a.pack('C*'))
    end
  end
end
