content: package org.catacombae.hfsexplorer.iphone;
  
  import java.util.Arrays;
  import java.util.HashMap;
  import java.nio.ByteBuffer;
  import javax.crypto.Cipher;
  import javax.crypto.spec.SecretKeySpec;
  import org.catacombae.dmgextractor.Util;
  import de.rtner.security.auth.spi.PBKDF2Engine;
  import de.rtner.security.auth.spi.PBKDF2Parameters;
  import djb.Curve25519;
  import java.security.MessageDigest;
  import java.io.ByteArrayOutputStream;
  
  public class Keybag {
      private static final int WRAP_DEVICE = 1;
      private static final int WRAP_PASSCODE = 2;
      public static enum Types {SYSTEM_KEYBAG, BACKUP_KEYBAG, ESCROW_KEYBAG, OTA_KEYBAG};
      private static final String[] KEYBAG_TAGS = {"VERS", "TYPE", "UUID", "HMCK", "WRAP", "SALT", "ITER"};
      private static final String[] CLASSKEY_TAGS = {"CLAS","WRAP","WPKY", "KTYP", "PBKY"};
      
      private HashMap<String, byte[]> attributes;
      private HashMap<byte[], HashMap<String, byte[]>> classKeys;
      public Boolean unlocked = false; // not sure if needed
  
      public Keybag(byte[] data)
      {
          this.attributes = new HashMap<String, byte[]>();
          this.classKeys = new HashMap<byte[], HashMap<String, byte[]>>();
          HashMap<String, byte[]> currentClassKey = new HashMap<String, byte[]>();
  
          for(int i=0; i + 8 < data.length; )
          {
              String tag = Util.toASCIIString(data, i, 4);
              int len = Util.readIntBE(data, i+4);
              byte[] value = Arrays.copyOfRange(data, i+8, i+8+len);
              if (tag.equals("UUID") && !this.attributes.containsKey("UUID"))
              {
                  this.attributes.put(tag, value);
              }
              else if (tag.equals("WRAP") && !this.attributes.containsKey("WRAP"))
              {
                  this.attributes.put(tag, value);
              }
              else if (tag.equals("UUID"))
              {
                  if (!currentClassKey.isEmpty())
                  {
                      classKeys.put(currentClassKey.get("CLAS"), currentClassKey);
                  }
  
                  currentClassKey = new HashMap<String, byte[]>();
                  currentClassKey.put("UUID", value);
              }
              else if (Arrays.asList(CLASSKEY_TAGS).contains(tag))
              {
                  currentClassKey.put(tag, value);
              }
              else
              {
                  this.attributes.put(tag, value);
              }
              i += 8 + len;
          }
      }
  
      public byte[] getSALT() {
          return this.attributes.get("SALT");
      }
  
      public int getITER() {
          return ByteBuffer.wrap(this.attributes.get("ITER")).getInt();
      }
  
      public Types getTYPE() {
          return Types.values()[ByteBuffer.wrap(this.attributes.get("TYPE")).getInt()];
      }
      
      public Boolean unlockWithPasscodeKey(byte[] passcodeKey) throws Exception {
          for (HashMap<String, byte[]> classKey : this.classKeys.values()) {
              if (!classKey.containsKey("WPKY"))
                  continue;
              
              byte[] k = classKey.get("WPKY");
              if ((ByteBuffer.wrap(classKey.get("WRAP")).getInt() & WRAP_PASSCODE) != 0) {
                  k = AESUnwrap(passcodeKey, classKey.get("WPKY"));
              }
              // this is not needed
              /*if ((wrap & WRAP_DEVICE) != 0) { 
                  //python code:
                  //if not self.deviceKey:
                  //    continue
                  //k = AESdecryptCBC(k, self.deviceKey)
                  
              }*/
              classKey.put("KEY", k);
          }
          this.unlocked = true;
          
          return true;
      }
      
      public byte[] getPasscodekeyFromPasscode(byte[] passcode) throws Exception {
          int iter = this.getTYPE() == Types.BACKUP_KEYBAG || this.getTYPE() == Types.OTA_KEYBAG ? this.getITER() : 1;
          PBKDF2Parameters p = new PBKDF2Parameters("HmacSHA1", null, this.getSALT(), iter);
          PBKDF2Engine e = new PBKDF2Engine(p);
          return e.deriveKey(passcode, 32);
      }
      
      public Boolean unlockBackupKeybagWithPasscode(byte[] passcode) throws Exception {
          if (this.getTYPE() == Types.BACKUP_KEYBAG || this.getTYPE() == Types.OTA_KEYBAG) {
              return false; // not a backup keybag!
          }
          return this.unlockWithPasscodeKey(this.getPasscodekeyFromPasscode(passcode));
      }
      
      private byte[] AESUnwrap(byte[] key, byte[] wrapped) throws Exception {
          Cipher cipher = Cipher.getInstance("AESWrap");
          cipher.init(Cipher.UNWRAP_MODE, new SecretKeySpec(key, "AES"));
          return cipher.unwrap(wrapped, "AES", Cipher.SECRET_KEY).getEncoded();
      }
      
      private byte[] unwrapCurve25519(byte[] persistentClass, byte[] persistentKey) throws Exception {
          assert persistentKey.length == 0x48;
          byte[] mySecret = this.classKeys.get(persistentClass).get("KEY");
          byte[] myPublic = this.classKeys.get(persistentClass).get("PBKY");
          byte[] hisPublic = Arrays.copyOfRange(persistentKey, 0, 32);
          byte[] shared = new byte[32];
          Curve25519.curve(shared, mySecret, hisPublic); // not really sure
          
          ByteArrayOutputStream stream = new ByteArrayOutputStream();
          stream.write(new byte[] { 0x00, 0x00, 0x00, 0x01 });
          stream.write(shared);
          stream.write(hisPublic);
          stream.write(myPublic);
          
          MessageDigest digest = MessageDigest.getInstance("SHA-256");
          byte[] md = digest.digest(stream.toByteArray());
          
          return AESUnwrap(md, hisPublic);
      }
  }
  