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