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             logger.debug("Refreshing {} ec pre keys, because only {} of {} pre keys remain",
  42                     serviceIdType,
  43                     preKeyCounts.getEcCount(),
  44                     ServiceConfig.PREKEY_MINIMUM_COUNT);
  45             refreshPreKeys(serviceIdType);
  46         }
  47         if (preKeyCounts.getKyberCount() < ServiceConfig.PREKEY_MINIMUM_COUNT) {
  48             logger.debug("Refreshing {} kyber pre keys, because only {} of {} pre keys remain",
  49                     serviceIdType,
  50                     preKeyCounts.getEcCount(),
  51                     ServiceConfig.PREKEY_MINIMUM_COUNT);
  52             refreshKyberPreKeys(serviceIdType);
  53         }
  54     }
  55
  56     private void refreshPreKeys(ServiceIdType serviceIdType) throws IOException {
  57         final var identityKeyPair = account.getIdentityKeyPair(serviceIdType);
  58         if (identityKeyPair == null) {
  59             return;
  60         }
  61         final var accountId = account.getAccountId(serviceIdType);
  62         if (accountId == null) {
  63             return;
  64         }
  65         try {
  66             refreshPreKeys(serviceIdType, identityKeyPair);
  67         } catch (Exception e) {
  68             logger.warn("Failed to store new pre keys, resetting preKey id offset", e);
  69             account.resetPreKeyOffsets(serviceIdType);
  70             refreshPreKeys(serviceIdType, identityKeyPair);
  71         }
  72     }
  73
  74     private void refreshPreKeys(
  75             final ServiceIdType serviceIdType, final IdentityKeyPair identityKeyPair
  76     ) throws IOException {
  77         final var oneTimePreKeys = generatePreKeys(serviceIdType);
  78         final var signedPreKeyRecord = generateSignedPreKey(serviceIdType, identityKeyPair);
  79
  80         final var preKeyUpload = new PreKeyUpload(serviceIdType,
  81                 identityKeyPair.getPublicKey(),
  82                 signedPreKeyRecord,
  83                 oneTimePreKeys,
  84                 null,
  85                 null);
  86         dependencies.getAccountManager().setPreKeys(preKeyUpload);
  87     }
  88
  89     private List<PreKeyRecord> generatePreKeys(ServiceIdType serviceIdType) {
  90         final var offset = account.getPreKeyIdOffset(serviceIdType);
  91
  92         var records = KeyUtils.generatePreKeyRecords(offset);
  93         account.addPreKeys(serviceIdType, records);
  94
  95         return records;
  96     }
  97
  98     private SignedPreKeyRecord generateSignedPreKey(ServiceIdType serviceIdType, IdentityKeyPair identityKeyPair) {
  99         final var signedPreKeyId = account.getNextSignedPreKeyId(serviceIdType);
 100
 101         var record = KeyUtils.generateSignedPreKeyRecord(signedPreKeyId, identityKeyPair);
 102         account.addSignedPreKey(serviceIdType, record);
 103
 104         return record;
 105     }
 106
 107     private void refreshKyberPreKeys(ServiceIdType serviceIdType) throws IOException {
 108         final var identityKeyPair = account.getIdentityKeyPair(serviceIdType);
 109         if (identityKeyPair == null) {
 110             return;
 111         }
 112         final var accountId = account.getAccountId(serviceIdType);
 113         if (accountId == null) {
 114             return;
 115         }
 116         try {
 117             refreshKyberPreKeys(serviceIdType, identityKeyPair);
 118         } catch (Exception e) {
 119             logger.warn("Failed to store new pre keys, resetting preKey id offset", e);
 120             account.resetKyberPreKeyOffsets(serviceIdType);
 121             refreshKyberPreKeys(serviceIdType, identityKeyPair);
 122         }
 123     }
 124
 125     private void refreshKyberPreKeys(
 126             final ServiceIdType serviceIdType, final IdentityKeyPair identityKeyPair
 127     ) throws IOException {
 128         final var oneTimePreKeys = generateKyberPreKeys(serviceIdType, identityKeyPair);
 129         final var lastResortPreKeyRecord = generateLastResortKyberPreKey(serviceIdType, identityKeyPair);
 130
 131         final var preKeyUpload = new PreKeyUpload(serviceIdType,
 132                 identityKeyPair.getPublicKey(),
 133                 null,
 134                 null,
 135                 lastResortPreKeyRecord,
 136                 oneTimePreKeys);
 137         dependencies.getAccountManager().setPreKeys(preKeyUpload);
 138     }
 139
 140     private List<KyberPreKeyRecord> generateKyberPreKeys(
 141             ServiceIdType serviceIdType, final IdentityKeyPair identityKeyPair
 142     ) {
 143         final var offset = account.getKyberPreKeyIdOffset(serviceIdType);
 144
 145         var records = KeyUtils.generateKyberPreKeyRecords(offset, identityKeyPair.getPrivateKey());
 146         account.addKyberPreKeys(serviceIdType, records);
 147
 148         return records;
 149     }
 150
 151     private KyberPreKeyRecord generateLastResortKyberPreKey(
 152             ServiceIdType serviceIdType, IdentityKeyPair identityKeyPair
 153     ) {
 154         final var signedPreKeyId = account.getKyberPreKeyIdOffset(serviceIdType);
 155
 156         var record = KeyUtils.generateKyberPreKeyRecord(signedPreKeyId, identityKeyPair.getPrivateKey());
 157         account.addLastResortKyberPreKey(serviceIdType, record);
 158
 159         return record;
 160     }
 161 }