michael@0: /* This Source Code Form is subject to the terms of the Mozilla Public
michael@0:  * License, v. 2.0. If a copy of the MPL was not distributed with this
michael@0:  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
michael@0: 
michael@0: package org.mozilla.gecko.sync;
michael@0: 
michael@0: import java.io.IOException;
michael@0: import java.io.UnsupportedEncodingException;
michael@0: import java.util.HashMap;
michael@0: import java.util.HashSet;
michael@0: import java.util.Map.Entry;
michael@0: import java.util.Set;
michael@0: 
michael@0: import org.json.simple.JSONArray;
michael@0: import org.json.simple.parser.ParseException;
michael@0: import org.mozilla.apache.commons.codec.binary.Base64;
michael@0: import org.mozilla.gecko.sync.crypto.CryptoException;
michael@0: import org.mozilla.gecko.sync.crypto.KeyBundle;
michael@0: 
michael@0: public class CollectionKeys {
michael@0:   private KeyBundle                  defaultKeyBundle     = null;
michael@0:   private final HashMap<String, KeyBundle> collectionKeyBundles = new HashMap<String, KeyBundle>();
michael@0: 
michael@0:   /**
michael@0:    * Randomly generate a basic CollectionKeys object.
michael@0:    * @throws CryptoException
michael@0:    */
michael@0:   public static CollectionKeys generateCollectionKeys() throws CryptoException {
michael@0:     CollectionKeys ck = new CollectionKeys();
michael@0:     ck.clear();
michael@0:     ck.defaultKeyBundle = KeyBundle.withRandomKeys();
michael@0:     // TODO: eventually we would like to keep per-collection keys, just generate
michael@0:     // new ones as appropriate.
michael@0:     return ck;
michael@0:   }
michael@0: 
michael@0:   public KeyBundle defaultKeyBundle() throws NoCollectionKeysSetException {
michael@0:     if (this.defaultKeyBundle == null) {
michael@0:       throw new NoCollectionKeysSetException();
michael@0:     }
michael@0:     return this.defaultKeyBundle;
michael@0:   }
michael@0: 
michael@0:   public boolean keyBundleForCollectionIsNotDefault(String collection) {
michael@0:     return collectionKeyBundles.containsKey(collection);
michael@0:   }
michael@0: 
michael@0:   public KeyBundle keyBundleForCollection(String collection)
michael@0:       throws NoCollectionKeysSetException {
michael@0:     if (this.defaultKeyBundle == null) {
michael@0:       throw new NoCollectionKeysSetException();
michael@0:     }
michael@0:     if (keyBundleForCollectionIsNotDefault(collection)) {
michael@0:       return collectionKeyBundles.get(collection);
michael@0:     }
michael@0:     return this.defaultKeyBundle;
michael@0:   }
michael@0: 
michael@0:   /**
michael@0:    * Take a pair of values in a JSON array, handing them off to KeyBundle to
michael@0:    * produce a usable keypair.
michael@0:    */
michael@0:   private static KeyBundle arrayToKeyBundle(JSONArray array) throws UnsupportedEncodingException {
michael@0:     String encKeyStr  = (String) array.get(0);
michael@0:     String hmacKeyStr = (String) array.get(1);
michael@0:     return KeyBundle.fromBase64EncodedKeys(encKeyStr, hmacKeyStr);
michael@0:   }
michael@0: 
michael@0:   @SuppressWarnings("unchecked")
michael@0:   private static JSONArray keyBundleToArray(KeyBundle bundle) {
michael@0:     // Generate JSON.
michael@0:     JSONArray keysArray = new JSONArray();
michael@0:     keysArray.add(new String(Base64.encodeBase64(bundle.getEncryptionKey())));
michael@0:     keysArray.add(new String(Base64.encodeBase64(bundle.getHMACKey())));
michael@0:     return keysArray;
michael@0:   }
michael@0: 
michael@0:   private ExtendedJSONObject asRecordContents() throws NoCollectionKeysSetException {
michael@0:     ExtendedJSONObject json = new ExtendedJSONObject();
michael@0:     json.put("id", "keys");
michael@0:     json.put("collection", "crypto");
michael@0:     json.put("default", keyBundleToArray(this.defaultKeyBundle()));
michael@0:     ExtendedJSONObject colls = new ExtendedJSONObject();
michael@0:     for (Entry<String, KeyBundle> collKey : collectionKeyBundles.entrySet()) {
michael@0:       colls.put(collKey.getKey(), keyBundleToArray(collKey.getValue()));
michael@0:     }
michael@0:     json.put("collections", colls);
michael@0:     return json;
michael@0:   }
michael@0: 
michael@0:   public CryptoRecord asCryptoRecord() throws NoCollectionKeysSetException {
michael@0:     ExtendedJSONObject payload = this.asRecordContents();
michael@0:     CryptoRecord record = new CryptoRecord(payload);
michael@0:     record.collection = "crypto";
michael@0:     record.guid       = "keys";
michael@0:     record.deleted    = false;
michael@0:     return record;
michael@0:   }
michael@0: 
michael@0:   /**
michael@0:    * Set my key bundle and collection keys with the given key bundle and data
michael@0:    * (possibly decrypted) from the given record.
michael@0:    *
michael@0:    * @param keys
michael@0:    *          A "crypto/keys" <code>CryptoRecord</code>, encrypted with
michael@0:    *          <code>syncKeyBundle</code> if <code>syncKeyBundle</code> is non-null.
michael@0:    * @param syncKeyBundle
michael@0:    *          If non-null, the sync key bundle to decrypt <code>keys</code> with.
michael@0:    */
michael@0:   public void setKeyPairsFromWBO(CryptoRecord keys, KeyBundle syncKeyBundle)
michael@0:       throws CryptoException, IOException, ParseException, NonObjectJSONException {
michael@0:     if (keys == null) {
michael@0:       throw new IllegalArgumentException("cannot set key pairs from null record");
michael@0:     }
michael@0:     if (syncKeyBundle != null) {
michael@0:       keys.keyBundle = syncKeyBundle;
michael@0:       keys.decrypt();
michael@0:     }
michael@0:     ExtendedJSONObject cleartext = keys.payload;
michael@0:     KeyBundle defaultKey = arrayToKeyBundle((JSONArray) cleartext.get("default"));
michael@0: 
michael@0:     ExtendedJSONObject collections = cleartext.getObject("collections");
michael@0:     HashMap<String, KeyBundle> collectionKeys = new HashMap<String, KeyBundle>();
michael@0:     for (Entry<String, Object> pair : collections.entrySet()) {
michael@0:       KeyBundle bundle = arrayToKeyBundle((JSONArray) pair.getValue());
michael@0:       collectionKeys.put(pair.getKey(), bundle);
michael@0:     }
michael@0: 
michael@0:     this.collectionKeyBundles.clear();
michael@0:     this.collectionKeyBundles.putAll(collectionKeys);
michael@0:     this.defaultKeyBundle     = defaultKey;
michael@0:   }
michael@0: 
michael@0:   public void setKeyBundleForCollection(String collection, KeyBundle keys) {
michael@0:     this.collectionKeyBundles.put(collection, keys);
michael@0:   }
michael@0: 
michael@0:   public void setDefaultKeyBundle(KeyBundle keys) {
michael@0:     this.defaultKeyBundle = keys;
michael@0:   }
michael@0: 
michael@0:   public void clear() {
michael@0:     this.defaultKeyBundle = null;
michael@0:     this.collectionKeyBundles.clear();
michael@0:   }
michael@0: 
michael@0:   /**
michael@0:    * Return set of collections where key is either missing from one collection
michael@0:    * or not the same in both collections.
michael@0:    * <p>
michael@0:    * Does not check for different default keys.
michael@0:    */
michael@0:   public static Set<String> differences(CollectionKeys a, CollectionKeys b) {
michael@0:     Set<String> differences = new HashSet<String>();
michael@0:     Set<String> collections = new HashSet<String>(a.collectionKeyBundles.keySet());
michael@0:     collections.addAll(b.collectionKeyBundles.keySet());
michael@0: 
michael@0:     // Iterate through one collection, collecting missing and differences.
michael@0:     for (String collection : collections) {
michael@0:       KeyBundle keyA;
michael@0:       KeyBundle keyB;
michael@0:       try {
michael@0:         keyA = a.keyBundleForCollection(collection); // Will return default key as appropriate.
michael@0:         keyB = b.keyBundleForCollection(collection); // Will return default key as appropriate.
michael@0:       } catch (NoCollectionKeysSetException e) {
michael@0:         differences.add(collection);
michael@0:         continue;
michael@0:       }
michael@0:       // keyA and keyB are not null at this point.
michael@0:       if (!keyA.equals(keyB)) {
michael@0:         differences.add(collection);
michael@0:       }
michael@0:     }
michael@0: 
michael@0:     return differences;
michael@0:   }
michael@0: 
michael@0:   @Override
michael@0:   public boolean equals(Object o) {
michael@0:     if (!(o instanceof CollectionKeys)) {
michael@0:       return false;
michael@0:     }
michael@0:     CollectionKeys other = (CollectionKeys) o;
michael@0:     try {
michael@0:       // It would be nice to use map equality here, but there can be map entries
michael@0:       // where the key is the default key that should compare equal to a missing
michael@0:       // map entry. Therefore, we always compute the set of differences.
michael@0:       return defaultKeyBundle().equals(other.defaultKeyBundle()) &&
michael@0:              CollectionKeys.differences(this, other).isEmpty();
michael@0:     } catch (NoCollectionKeysSetException e) {
michael@0:       // If either default key bundle is not set, we'll say the bundles are not equal.
michael@0:       return false;
michael@0:     }
michael@0:   }
michael@0: }