nmode's Git Repositories - signal-cli/blob - lib/src/main/java/org/asamk/signal/manager/helper/PreKeyHelper.java

   1 package org.asamk.signal.manager.helper;
   2
   3 import org.asamk.signal.manager.config.ServiceConfig;
   4 import org.asamk.signal.manager.internal.SignalDependencies;
   5 import org.asamk.signal.manager.storage.SignalAccount;
   6 import org.asamk.signal.manager.util.KeyUtils;
   7 import org.signal.libsignal.protocol.IdentityKeyPair;
   8 import org.signal.libsignal.protocol.state.PreKeyRecord;
   9 import org.signal.libsignal.protocol.state.SignedPreKeyRecord;
  10 import org.slf4j.Logger;
  11 import org.slf4j.LoggerFactory;
  12 import org.whispersystems.signalservice.api.account.PreKeyUpload;
  13 import org.whispersystems.signalservice.api.push.ServiceIdType;
  14
  15 import java.io.IOException;
  16 import java.util.List;
  17
  18 public class PreKeyHelper {
  19
  20     private final static Logger logger = LoggerFactory.getLogger(PreKeyHelper.class);
  21
  22     private final SignalAccount account;
  23     private final SignalDependencies dependencies;
  24
  25     public PreKeyHelper(
  26             final SignalAccount account, final SignalDependencies dependencies
  27     ) {
  28         this.account = account;
  29         this.dependencies = dependencies;
  30     }
  31
  32     public void refreshPreKeysIfNecessary() throws IOException {
  33         refreshPreKeysIfNecessary(ServiceIdType.ACI);
  34         refreshPreKeysIfNecessary(ServiceIdType.PNI);
  35     }
  36
  37     public void refreshPreKeysIfNecessary(ServiceIdType serviceIdType) throws IOException {
  38         final var preKeyCounts = dependencies.getAccountManager().getPreKeyCounts(serviceIdType);
  39         if (preKeyCounts.getEcCount() < ServiceConfig.PREKEY_MINIMUM_COUNT) {
  40             refreshPreKeys(serviceIdType);
  41         }
  42         // TODO kyber pre keys
  43     }
  44
  45     public void refreshPreKeys() throws IOException {
  46         refreshPreKeys(ServiceIdType.ACI);
  47         refreshPreKeys(ServiceIdType.PNI);
  48     }
  49
  50     public void refreshPreKeys(ServiceIdType serviceIdType) throws IOException {
  51         final var identityKeyPair = account.getIdentityKeyPair(serviceIdType);
  52         if (identityKeyPair == null) {
  53             return;
  54         }
  55         final var accountId = account.getAccountId(serviceIdType);
  56         if (accountId == null) {
  57             return;
  58         }
  59         try {
  60             refreshPreKeys(serviceIdType, identityKeyPair);
  61         } catch (Exception e) {
  62             logger.warn("Failed to store new pre keys, resetting preKey id offset", e);
  63             account.resetPreKeyOffsets(serviceIdType);
  64             refreshPreKeys(serviceIdType, identityKeyPair);
  65         }
  66     }
  67
  68     private void refreshPreKeys(
  69             final ServiceIdType serviceIdType, final IdentityKeyPair identityKeyPair
  70     ) throws IOException {
  71         final var oneTimePreKeys = generatePreKeys(serviceIdType);
  72         final var signedPreKeyRecord = generateSignedPreKey(serviceIdType, identityKeyPair);
  73
  74         final var preKeyUpload = new PreKeyUpload(serviceIdType,
  75                 identityKeyPair.getPublicKey(),
  76                 signedPreKeyRecord,
  77                 oneTimePreKeys,
  78                 null,
  79                 null);
  80         dependencies.getAccountManager().setPreKeys(preKeyUpload);
  81     }
  82
  83     private List<PreKeyRecord> generatePreKeys(ServiceIdType serviceIdType) {
  84         final var offset = account.getPreKeyIdOffset(serviceIdType);
  85
  86         var records = KeyUtils.generatePreKeyRecords(offset, ServiceConfig.PREKEY_BATCH_SIZE);
  87         account.addPreKeys(serviceIdType, records);
  88
  89         return records;
  90     }
  91
  92     private SignedPreKeyRecord generateSignedPreKey(ServiceIdType serviceIdType, IdentityKeyPair identityKeyPair) {
  93         final var signedPreKeyId = account.getNextSignedPreKeyId(serviceIdType);
  94
  95         var record = KeyUtils.generateSignedPreKeyRecord(identityKeyPair, signedPreKeyId);
  96         account.addSignedPreKey(serviceIdType, record);
  97
  98         return record;
  99     }
 100 }