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

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