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.KyberPreKeyRecord;
   9 import org.signal.libsignal.protocol.state.PreKeyRecord;
  10 import org.signal.libsignal.protocol.state.SignedPreKeyRecord;
  11 import org.slf4j.Logger;
  12 import org.slf4j.LoggerFactory;
  13 import org.whispersystems.signalservice.api.account.PreKeyUpload;
  14 import org.whispersystems.signalservice.api.push.ServiceIdType;
  15
  16 import java.io.IOException;
  17 import java.util.List;
  18
  19 public class PreKeyHelper {
  20
  21     private final static Logger logger = LoggerFactory.getLogger(PreKeyHelper.class);
  22
  23     private final SignalAccount account;
  24     private final SignalDependencies dependencies;
  25
  26     public PreKeyHelper(
  27             final SignalAccount account, final SignalDependencies dependencies
  28     ) {
  29         this.account = account;
  30         this.dependencies = dependencies;
  31     }
  32
  33     public void refreshPreKeysIfNecessary() throws IOException {
  34         refreshPreKeysIfNecessary(ServiceIdType.ACI);
  35         refreshPreKeysIfNecessary(ServiceIdType.PNI);
  36     }
  37
  38     public void refreshPreKeysIfNecessary(ServiceIdType serviceIdType) throws IOException {
  39         final var preKeyCounts = dependencies.getAccountManager().getPreKeyCounts(serviceIdType);
  40         if (preKeyCounts.getEcCount() < ServiceConfig.PREKEY_MINIMUM_COUNT) {
  41             refreshPreKeys(serviceIdType);
  42         }
  43         if (preKeyCounts.getKyberCount() < ServiceConfig.PREKEY_MINIMUM_COUNT) {
  44             refreshKyberPreKeys(serviceIdType);
  45         }
  46     }
  47
  48     public void refreshPreKeys() throws IOException {
  49         refreshPreKeys(ServiceIdType.ACI);
  50         refreshPreKeys(ServiceIdType.PNI);
  51     }
  52
  53     private void refreshPreKeys(ServiceIdType serviceIdType) throws IOException {
  54         final var identityKeyPair = account.getIdentityKeyPair(serviceIdType);
  55         if (identityKeyPair == null) {
  56             return;
  57         }
  58         final var accountId = account.getAccountId(serviceIdType);
  59         if (accountId == null) {
  60             return;
  61         }
  62         try {
  63             refreshPreKeys(serviceIdType, identityKeyPair);
  64         } catch (Exception e) {
  65             logger.warn("Failed to store new pre keys, resetting preKey id offset", e);
  66             account.resetPreKeyOffsets(serviceIdType);
  67             refreshPreKeys(serviceIdType, identityKeyPair);
  68         }
  69     }
  70
  71     private void refreshPreKeys(
  72             final ServiceIdType serviceIdType, final IdentityKeyPair identityKeyPair
  73     ) throws IOException {
  74         final var oneTimePreKeys = generatePreKeys(serviceIdType);
  75         final var signedPreKeyRecord = generateSignedPreKey(serviceIdType, identityKeyPair);
  76
  77         final var preKeyUpload = new PreKeyUpload(serviceIdType,
  78                 identityKeyPair.getPublicKey(),
  79                 signedPreKeyRecord,
  80                 oneTimePreKeys,
  81                 null,
  82                 null);
  83         dependencies.getAccountManager().setPreKeys(preKeyUpload);
  84     }
  85
  86     private List<PreKeyRecord> generatePreKeys(ServiceIdType serviceIdType) {
  87         final var offset = account.getPreKeyIdOffset(serviceIdType);
  88
  89         var records = KeyUtils.generatePreKeyRecords(offset, ServiceConfig.PREKEY_BATCH_SIZE);
  90         account.addPreKeys(serviceIdType, records);
  91
  92         return records;
  93     }
  94
  95     private SignedPreKeyRecord generateSignedPreKey(ServiceIdType serviceIdType, IdentityKeyPair identityKeyPair) {
  96         final var signedPreKeyId = account.getNextSignedPreKeyId(serviceIdType);
  97
  98         var record = KeyUtils.generateSignedPreKeyRecord(identityKeyPair, signedPreKeyId);
  99         account.addSignedPreKey(serviceIdType, record);
 100
 101         return record;
 102     }
 103
 104     private void refreshKyberPreKeys(ServiceIdType serviceIdType) throws IOException {
 105         final var identityKeyPair = account.getIdentityKeyPair(serviceIdType);
 106         if (identityKeyPair == null) {
 107             return;
 108         }
 109         final var accountId = account.getAccountId(serviceIdType);
 110         if (accountId == null) {
 111             return;
 112         }
 113         try {
 114             refreshKyberPreKeys(serviceIdType, identityKeyPair);
 115         } catch (Exception e) {
 116             logger.warn("Failed to store new pre keys, resetting preKey id offset", e);
 117             account.resetKyberPreKeyOffsets(serviceIdType);
 118             refreshKyberPreKeys(serviceIdType, identityKeyPair);
 119         }
 120     }
 121
 122     private void refreshKyberPreKeys(
 123             final ServiceIdType serviceIdType, final IdentityKeyPair identityKeyPair
 124     ) throws IOException {
 125         final var oneTimePreKeys = generateKyberPreKeys(serviceIdType, identityKeyPair);
 126         final var lastResortPreKeyRecord = generateLastResortKyberPreKey(serviceIdType, identityKeyPair);
 127
 128         final var preKeyUpload = new PreKeyUpload(serviceIdType,
 129                 identityKeyPair.getPublicKey(),
 130                 null,
 131                 null,
 132                 lastResortPreKeyRecord,
 133                 oneTimePreKeys);
 134         dependencies.getAccountManager().setPreKeys(preKeyUpload);
 135     }
 136
 137     private List<KyberPreKeyRecord> generateKyberPreKeys(
 138             ServiceIdType serviceIdType, final IdentityKeyPair identityKeyPair
 139     ) {
 140         final var offset = account.getKyberPreKeyIdOffset(serviceIdType);
 141
 142         var records = KeyUtils.generateKyberPreKeyRecords(offset,
 143                 ServiceConfig.PREKEY_BATCH_SIZE,
 144                 identityKeyPair.getPrivateKey());
 145         account.addKyberPreKeys(serviceIdType, records);
 146
 147         return records;
 148     }
 149
 150     private KyberPreKeyRecord generateLastResortKyberPreKey(
 151             ServiceIdType serviceIdType, IdentityKeyPair identityKeyPair
 152     ) {
 153         final var signedPreKeyId = account.getKyberPreKeyIdOffset(serviceIdType);
 154
 155         var record = KeyUtils.generateKyberPreKeyRecord(signedPreKeyId, identityKeyPair.getPrivateKey());
 156         account.addLastResortKyberPreKey(serviceIdType, record);
 157
 158         return record;
 159     }
 160 }