require 'ipcrypt/exceptions'
require 'securerandom'
module IPCrypt
  class Engine
    attr_reader :default_key
    IPv4 = /^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$/

    def initialize(*ips)
      @default_key = SecureRandom.random_bytes(16)
      @ips = ips.flatten.map do |ip|
        raise InvalidIPv4Error.new(ip) unless IPv4.match? ip
        ip
      end
    end

    def encrypt(key = @default_key)
      raise InvalidKeyTypeError.new(key, key.class) unless key.is_a? String

      bytes = key.bytesize
      raise InvalidKeyBytesError.new(key, bytes) unless bytes == 16

      encrypted = @ips.map do |ip|
        k     = key.chars.map {|c| c.unpack('<C').first }
        state = ip.split(?.).map &:to_i

        state = xor4 state, k[0...4]
        state = permute_fwd state
        state = xor4 state, k[4...8]
        state = permute_fwd state
        state = xor4 state, k[8...12]
        state = permute_fwd state
        state = xor4 state, k[12...16]

        state * ?.
      end

      encrypted.size == 1 ? encrypted.first : encrypted
    end

    def decrypt(key)
      raise InvalidKeyTypeError.new(key, key.class) unless key.is_a? String

      bytes = key.bytesize
      raise InvalidKeyBytesError.new(key, bytes) unless bytes == 16

      decrypted = @ips.map do |ip|
        k     = key.chars.map {|c| c.unpack('<C').first }
        state = ip.split(?.).map &:to_i

        state = xor4 state, k[12...16]
        state = permute_bwd state
        state = xor4 state, k[8...12]
        state = permute_bwd state
        state = xor4 state, k[4...8]
        state = permute_bwd state
        state = xor4 state, k[0...4]

        state * ?.
      end

      decrypted.size == 1 ? decrypted.first : decrypted
    end

    private

    def rotl(b, r)
      ((b << r) & 0xff) | (b >> (8 - r))
    end

    def xor4(x, y)
      (0..3).map {|i| (x[i] ^ y[i]) & 0xff }
    end

    def permute_fwd(state)
      b0, b1, b2, b3 = state
      b0 += b1
      b2 += b3
      b0 &= 0xff
      b2 &= 0xff
      b1 = rotl b1, 2
      b3 = rotl b3, 5
      b1 ^= b0
      b3 ^= b2
      b0 = rotl b0, 4
      b0 += b3
      b2 += b1
      b0 &= 0xff
      b2 &= 0xff
      b1 = rotl b1, 3
      b3 = rotl b3, 7
      b1 ^= b2
      b3 ^= b0
      b2 = rotl b2, 4
      [b0, b1, b2, b3]
    end

    def permute_bwd(state)
      b0, b1, b2, b3 = state
      b2 = rotl b2, 4
      b1 ^= b2
      b3 ^= b0
      b1 = rotl b1, 5
      b3 = rotl b3, 1
      b0 -= b3
      b2 -= b1
      b0 &= 0xff
      b2 &= 0xff
      b0 = rotl b0, 4
      b1 ^= b0
      b3 ^= b2
      b1 = rotl b1, 6
      b3 = rotl b3, 3
      b0 -= b1
      b2 -= b3
      b0 &= 0xff
      b2 &= 0xff
      [b0, b1, b2, b3]
    end
  end
end