package org.asamk.signal.storage.protocol; import com.fasterxml.jackson.core.JsonGenerator; import com.fasterxml.jackson.core.JsonParser; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.*; import org.whispersystems.libsignal.SignalProtocolAddress; import org.whispersystems.libsignal.state.SessionRecord; import org.whispersystems.libsignal.state.SessionStore; import org.whispersystems.signalservice.internal.util.Base64; import java.io.IOException; import java.util.*; class JsonSessionStore implements SessionStore { private final Map<SignalProtocolAddress, byte[]> sessions = new HashMap<>(); public JsonSessionStore() { } public void addSessions(Map<SignalProtocolAddress, byte[]> sessions) { this.sessions.putAll(sessions); } @Override public synchronized SessionRecord loadSession(SignalProtocolAddress remoteAddress) { try { if (containsSession(remoteAddress)) { return new SessionRecord(sessions.get(remoteAddress)); } else { return new SessionRecord(); } } catch (IOException e) { throw new AssertionError(e); } } @Override public synchronized List<Integer> getSubDeviceSessions(String name) { List<Integer> deviceIds = new LinkedList<>(); for (SignalProtocolAddress key : sessions.keySet()) { if (key.getName().equals(name) && key.getDeviceId() != 1) { deviceIds.add(key.getDeviceId()); } } return deviceIds; } @Override public synchronized void storeSession(SignalProtocolAddress address, SessionRecord record) { sessions.put(address, record.serialize()); } @Override public synchronized boolean containsSession(SignalProtocolAddress address) { return sessions.containsKey(address); } @Override public synchronized void deleteSession(SignalProtocolAddress address) { sessions.remove(address); } @Override public synchronized void deleteAllSessions(String name) { for (SignalProtocolAddress key : new ArrayList<>(sessions.keySet())) { if (key.getName().equals(name)) { sessions.remove(key); } } } public static class JsonSessionStoreDeserializer extends JsonDeserializer<JsonSessionStore> { @Override public JsonSessionStore deserialize(JsonParser jsonParser, DeserializationContext deserializationContext) throws IOException, JsonProcessingException { JsonNode node = jsonParser.getCodec().readTree(jsonParser); Map<SignalProtocolAddress, byte[]> sessionMap = new HashMap<>(); if (node.isArray()) { for (JsonNode session : node) { String sessionName = session.get("name").asText(); try { sessionMap.put(new SignalProtocolAddress(sessionName, session.get("deviceId").asInt()), Base64.decode(session.get("record").asText())); } catch (IOException e) { System.out.println(String.format("Error while decoding session for: %s", sessionName)); } } } JsonSessionStore sessionStore = new JsonSessionStore(); sessionStore.addSessions(sessionMap); return sessionStore; } } public static class JsonPreKeyStoreSerializer extends JsonSerializer<JsonSessionStore> { @Override public void serialize(JsonSessionStore jsonSessionStore, JsonGenerator json, SerializerProvider serializerProvider) throws IOException, JsonProcessingException { json.writeStartArray(); for (Map.Entry<SignalProtocolAddress, byte[]> preKey : jsonSessionStore.sessions.entrySet()) { json.writeStartObject(); json.writeStringField("name", preKey.getKey().getName()); json.writeNumberField("deviceId", preKey.getKey().getDeviceId()); json.writeStringField("record", Base64.encodeBytes(preKey.getValue())); json.writeEndObject(); } json.writeEndArray(); } } }