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