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