X-Git-Url: https://git.nmode.ca/signal-cli/blobdiff_plain/a52f6a6657585fbad5afa4c57ce37752118317e9..425626ef9475cbc90ef8ada95dee172389baf521:/src/main/java/org/asamk/signal/manager/helper/PinHelper.java

diff --git a/src/main/java/org/asamk/signal/manager/helper/PinHelper.java b/src/main/java/org/asamk/signal/manager/helper/PinHelper.java
new file mode 100644
index 00000000..47ee6b40
--- /dev/null
+++ b/src/main/java/org/asamk/signal/manager/helper/PinHelper.java
@@ -0,0 +1,90 @@
+package org.asamk.signal.manager.helper;
+
+import org.asamk.signal.manager.util.PinHashing;
+import org.whispersystems.signalservice.api.KbsPinData;
+import org.whispersystems.signalservice.api.KeyBackupService;
+import org.whispersystems.signalservice.api.KeyBackupServicePinException;
+import org.whispersystems.signalservice.api.KeyBackupSystemNoDataException;
+import org.whispersystems.signalservice.api.kbs.HashedPin;
+import org.whispersystems.signalservice.api.kbs.MasterKey;
+import org.whispersystems.signalservice.internal.contacts.crypto.UnauthenticatedResponseException;
+import org.whispersystems.signalservice.internal.contacts.entities.TokenResponse;
+import org.whispersystems.signalservice.internal.push.LockedException;
+
+import java.io.IOException;
+
+public class PinHelper {
+
+    private final KeyBackupService keyBackupService;
+
+    public PinHelper(final KeyBackupService keyBackupService) {
+        this.keyBackupService = keyBackupService;
+    }
+
+    public void setRegistrationLockPin(
+            String pin, MasterKey masterKey
+    ) throws IOException, UnauthenticatedResponseException {
+        final KeyBackupService.PinChangeSession pinChangeSession = keyBackupService.newPinChangeSession();
+        final HashedPin hashedPin = PinHashing.hashPin(pin, pinChangeSession);
+
+        pinChangeSession.setPin(hashedPin, masterKey);
+        pinChangeSession.enableRegistrationLock(masterKey);
+    }
+
+    public void removeRegistrationLockPin() throws IOException, UnauthenticatedResponseException {
+        final KeyBackupService.PinChangeSession pinChangeSession = keyBackupService.newPinChangeSession();
+        pinChangeSession.removePin();
+    }
+
+    public KbsPinData getRegistrationLockData(
+            String pin,
+            LockedException e
+    ) throws IOException, KeyBackupSystemNoDataException, KeyBackupServicePinException {
+        String basicStorageCredentials = e.getBasicStorageCredentials();
+        if (basicStorageCredentials == null) {
+            return null;
+        }
+
+        return getRegistrationLockData(pin, basicStorageCredentials);
+    }
+
+    private KbsPinData getRegistrationLockData(
+            String pin,
+            String basicStorageCredentials
+    ) throws IOException, KeyBackupSystemNoDataException, KeyBackupServicePinException {
+        TokenResponse tokenResponse = keyBackupService.getToken(basicStorageCredentials);
+        if (tokenResponse == null || tokenResponse.getTries() == 0) {
+            throw new IOException("KBS Account locked");
+        }
+
+        KbsPinData registrationLockData = restoreMasterKey(pin, basicStorageCredentials, tokenResponse);
+        if (registrationLockData == null) {
+            throw new AssertionError("Failed to restore master key");
+        }
+        return registrationLockData;
+    }
+
+    private KbsPinData restoreMasterKey(
+            String pin, String basicStorageCredentials, TokenResponse tokenResponse
+    ) throws IOException, KeyBackupSystemNoDataException, KeyBackupServicePinException {
+        if (pin == null) return null;
+
+        if (basicStorageCredentials == null) {
+            throw new AssertionError("Cannot restore KBS key, no storage credentials supplied");
+        }
+
+        KeyBackupService.RestoreSession session = keyBackupService.newRegistrationSession(basicStorageCredentials,
+                tokenResponse);
+
+        try {
+            HashedPin hashedPin = PinHashing.hashPin(pin, session);
+            KbsPinData kbsData = session.restorePin(hashedPin);
+            if (kbsData == null) {
+                throw new AssertionError("Null not expected");
+            }
+            return kbsData;
+        } catch (UnauthenticatedResponseException e) {
+            throw new IOException(e);
+        }
+    }
+}