nmode's Git Repositories - signal-cli/blob - src/main/java/org/asamk/signal/storage/protocol/JsonPreKeyStore.java

   1 package org.asamk.signal.storage.protocol;
   2
   3 import com.fasterxml.jackson.core.JsonGenerator;
   4 import com.fasterxml.jackson.core.JsonParser;
   5 import com.fasterxml.jackson.databind.DeserializationContext;
   6 import com.fasterxml.jackson.databind.JsonDeserializer;
   7 import com.fasterxml.jackson.databind.JsonNode;
   8 import com.fasterxml.jackson.databind.JsonSerializer;
   9 import com.fasterxml.jackson.databind.SerializerProvider;
  10
  11 import org.whispersystems.libsignal.InvalidKeyIdException;
  12 import org.whispersystems.libsignal.state.PreKeyRecord;
  13 import org.whispersystems.libsignal.state.PreKeyStore;
  14 import org.whispersystems.util.Base64;
  15
  16 import java.io.IOException;
  17 import java.util.HashMap;
  18 import java.util.Map;
  19
  20 class JsonPreKeyStore implements PreKeyStore {
  21
  22     private final Map<Integer, byte[]> store = new HashMap<>();
  23
  24     public JsonPreKeyStore() {
  25
  26     }
  27
  28     private void addPreKeys(Map<Integer, byte[]> preKeys) {
  29         store.putAll(preKeys);
  30     }
  31
  32     @Override
  33     public PreKeyRecord loadPreKey(int preKeyId) throws InvalidKeyIdException {
  34         try {
  35             if (!store.containsKey(preKeyId)) {
  36                 throw new InvalidKeyIdException("No such prekeyrecord!");
  37             }
  38
  39             return new PreKeyRecord(store.get(preKeyId));
  40         } catch (IOException e) {
  41             throw new AssertionError(e);
  42         }
  43     }
  44
  45     @Override
  46     public void storePreKey(int preKeyId, PreKeyRecord record) {
  47         store.put(preKeyId, record.serialize());
  48     }
  49
  50     @Override
  51     public boolean containsPreKey(int preKeyId) {
  52         return store.containsKey(preKeyId);
  53     }
  54
  55     @Override
  56     public void removePreKey(int preKeyId) {
  57         store.remove(preKeyId);
  58     }
  59
  60     public static class JsonPreKeyStoreDeserializer extends JsonDeserializer<JsonPreKeyStore> {
  61
  62         @Override
  63         public JsonPreKeyStore deserialize(JsonParser jsonParser, DeserializationContext deserializationContext) throws IOException {
  64             JsonNode node = jsonParser.getCodec().readTree(jsonParser);
  65
  66             Map<Integer, byte[]> preKeyMap = new HashMap<>();
  67             if (node.isArray()) {
  68                 for (JsonNode preKey : node) {
  69                     Integer preKeyId = preKey.get("id").asInt();
  70                     try {
  71                         preKeyMap.put(preKeyId, Base64.decode(preKey.get("record").asText()));
  72                     } catch (IOException e) {
  73                         System.out.println(String.format("Error while decoding prekey for: %s", preKeyId));
  74                     }
  75                 }
  76             }
  77
  78             JsonPreKeyStore keyStore = new JsonPreKeyStore();
  79             keyStore.addPreKeys(preKeyMap);
  80
  81             return keyStore;
  82
  83         }
  84     }
  85
  86     public static class JsonPreKeyStoreSerializer extends JsonSerializer<JsonPreKeyStore> {
  87
  88         @Override
  89         public void serialize(JsonPreKeyStore jsonPreKeyStore, JsonGenerator json, SerializerProvider serializerProvider) throws IOException {
  90             json.writeStartArray();
  91             for (Map.Entry<Integer, byte[]> preKey : jsonPreKeyStore.store.entrySet()) {
  92                 json.writeStartObject();
  93                 json.writeNumberField("id", preKey.getKey());
  94                 json.writeStringField("record", Base64.encodeBytes(preKey.getValue()));
  95                 json.writeEndObject();
  96             }
  97             json.writeEndArray();
  98         }
  99     }
 100 }