nmode's Git Repositories - signal-cli/blob - src/main/java/org/asamk/signal/storage/protocol/JsonSignedPreKeyStore.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.SignedPreKeyRecord;
  11 import org.whispersystems.libsignal.state.SignedPreKeyStore;
  12
  13 import java.io.IOException;
  14 import java.util.HashMap;
  15 import java.util.LinkedList;
  16 import java.util.List;
  17 import java.util.Map;
  18
  19 class JsonSignedPreKeyStore implements SignedPreKeyStore {
  20
  21     private final Map<Integer, byte[]> store = new HashMap<>();
  22
  23     public JsonSignedPreKeyStore() {
  24
  25     }
  26
  27
  28     public void addSignedPreKeys(Map<Integer, byte[]> preKeys) {
  29         store.putAll(preKeys);
  30     }
  31
  32     @Override
  33     public SignedPreKeyRecord loadSignedPreKey(int signedPreKeyId) throws InvalidKeyIdException {
  34         try {
  35             if (!store.containsKey(signedPreKeyId)) {
  36                 throw new InvalidKeyIdException("No such signedprekeyrecord! " + signedPreKeyId);
  37             }
  38
  39             return new SignedPreKeyRecord(store.get(signedPreKeyId));
  40         } catch (IOException e) {
  41             throw new AssertionError(e);
  42         }
  43     }
  44
  45     @Override
  46     public List<SignedPreKeyRecord> loadSignedPreKeys() {
  47         try {
  48             List<SignedPreKeyRecord> results = new LinkedList<>();
  49
  50             for (byte[] serialized : store.values()) {
  51                 results.add(new SignedPreKeyRecord(serialized));
  52             }
  53
  54             return results;
  55         } catch (IOException e) {
  56             throw new AssertionError(e);
  57         }
  58     }
  59
  60     @Override
  61     public void storeSignedPreKey(int signedPreKeyId, SignedPreKeyRecord record) {
  62         store.put(signedPreKeyId, record.serialize());
  63     }
  64
  65     @Override
  66     public boolean containsSignedPreKey(int signedPreKeyId) {
  67         return store.containsKey(signedPreKeyId);
  68     }
  69
  70     @Override
  71     public void removeSignedPreKey(int signedPreKeyId) {
  72         store.remove(signedPreKeyId);
  73     }
  74
  75     public static class JsonSignedPreKeyStoreDeserializer extends JsonDeserializer<JsonSignedPreKeyStore> {
  76
  77         @Override
  78         public JsonSignedPreKeyStore deserialize(JsonParser jsonParser, DeserializationContext deserializationContext) throws IOException, JsonProcessingException {
  79             JsonNode node = jsonParser.getCodec().readTree(jsonParser);
  80
  81
  82             Map<Integer, byte[]> preKeyMap = new HashMap<>();
  83             if (node.isArray()) {
  84                 for (JsonNode preKey : node) {
  85                     Integer preKeyId = preKey.get("id").asInt();
  86                     try {
  87                         preKeyMap.put(preKeyId, Base64.decode(preKey.get("record").asText()));
  88                     } catch (IOException e) {
  89                         System.out.println(String.format("Error while decoding prekey for: %s", preKeyId));
  90                     }
  91                 }
  92             }
  93
  94             JsonSignedPreKeyStore keyStore = new JsonSignedPreKeyStore();
  95             keyStore.addSignedPreKeys(preKeyMap);
  96
  97             return keyStore;
  98
  99         }
 100     }
 101
 102     public static class JsonSignedPreKeyStoreSerializer extends JsonSerializer<JsonSignedPreKeyStore> {
 103
 104         @Override
 105         public void serialize(JsonSignedPreKeyStore jsonPreKeyStore, JsonGenerator json, SerializerProvider serializerProvider) throws IOException, JsonProcessingException {
 106             json.writeStartArray();
 107             for (Map.Entry<Integer, byte[]> signedPreKey : jsonPreKeyStore.store.entrySet()) {
 108                 json.writeStartObject();
 109                 json.writeNumberField("id", signedPreKey.getKey());
 110                 json.writeStringField("record", Base64.encodeBytes(signedPreKey.getValue()));
 111                 json.writeEndObject();
 112             }
 113             json.writeEndArray();
 114         }
 115     }
 116 }