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