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