001/**
002 *
003 * Copyright 2017 Paul Schaub
004 *
005 * Licensed under the Apache License, Version 2.0 (the "License");
006 * you may not use this file except in compliance with the License.
007 * You may obtain a copy of the License at
008 *
009 *     http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the License for the specific language governing permissions and
015 * limitations under the License.
016 */
017package org.jivesoftware.smackx.omemo.util;
018
019import static org.jivesoftware.smackx.omemo.util.OmemoConstants.Crypto.CIPHERMODE;
020import static org.jivesoftware.smackx.omemo.util.OmemoConstants.Crypto.KEYLENGTH;
021import static org.jivesoftware.smackx.omemo.util.OmemoConstants.Crypto.KEYTYPE;
022import static org.jivesoftware.smackx.omemo.util.OmemoConstants.Crypto.PROVIDER;
023
024import java.io.UnsupportedEncodingException;
025import java.security.InvalidAlgorithmParameterException;
026import java.security.InvalidKeyException;
027import java.security.NoSuchAlgorithmException;
028import java.security.NoSuchProviderException;
029import java.security.SecureRandom;
030import java.util.ArrayList;
031
032import javax.crypto.BadPaddingException;
033import javax.crypto.Cipher;
034import javax.crypto.IllegalBlockSizeException;
035import javax.crypto.KeyGenerator;
036import javax.crypto.NoSuchPaddingException;
037import javax.crypto.SecretKey;
038import javax.crypto.spec.IvParameterSpec;
039import javax.crypto.spec.SecretKeySpec;
040
041import org.jivesoftware.smack.util.StringUtils;
042
043import org.jivesoftware.smackx.omemo.OmemoManager;
044import org.jivesoftware.smackx.omemo.OmemoStore;
045import org.jivesoftware.smackx.omemo.element.OmemoVAxolotlElement;
046import org.jivesoftware.smackx.omemo.exceptions.CorruptedOmemoKeyException;
047import org.jivesoftware.smackx.omemo.exceptions.CryptoFailedException;
048import org.jivesoftware.smackx.omemo.exceptions.UndecidedOmemoIdentityException;
049import org.jivesoftware.smackx.omemo.internal.CiphertextTuple;
050import org.jivesoftware.smackx.omemo.internal.OmemoDevice;
051import org.jivesoftware.smackx.omemo.internal.OmemoSession;
052
053/**
054 * Class used to build OMEMO messages.
055 *
056 * @param <T_IdKeyPair> IdentityKeyPair class
057 * @param <T_IdKey>     IdentityKey class
058 * @param <T_PreKey>    PreKey class
059 * @param <T_SigPreKey> SignedPreKey class
060 * @param <T_Sess>      Session class
061 * @param <T_Addr>      Address class
062 * @param <T_ECPub>     Elliptic Curve PublicKey class
063 * @param <T_Bundle>    Bundle class
064 * @param <T_Ciph>      Cipher class
065 * @author Paul Schaub
066 */
067public class OmemoMessageBuilder<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> {
068    private final OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> omemoStore;
069    private final OmemoManager omemoManager;
070
071    private byte[] messageKey = generateKey();
072    private byte[] initializationVector = generateIv();
073
074    private byte[] ciphertextMessage;
075    private final ArrayList<OmemoVAxolotlElement.OmemoHeader.Key> keys = new ArrayList<>();
076
077    /**
078     * Create a OmemoMessageBuilder.
079     *
080     * @param omemoManager      OmemoManager of our device.
081     * @param omemoStore        OmemoStore.
082     * @param aesKey            AES key that will be transported to the recipient. This is used eg. to encrypt the body.
083     * @param iv                IV
084     * @throws NoSuchPaddingException
085     * @throws BadPaddingException
086     * @throws InvalidKeyException
087     * @throws NoSuchAlgorithmException
088     * @throws IllegalBlockSizeException
089     * @throws UnsupportedEncodingException
090     * @throws NoSuchProviderException
091     * @throws InvalidAlgorithmParameterException
092     */
093    public OmemoMessageBuilder(OmemoManager omemoManager,
094                               OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> omemoStore,
095                               byte[] aesKey, byte[] iv)
096            throws NoSuchPaddingException, BadPaddingException, InvalidKeyException, NoSuchAlgorithmException, IllegalBlockSizeException,
097            UnsupportedEncodingException, NoSuchProviderException, InvalidAlgorithmParameterException {
098        this.omemoStore = omemoStore;
099        this.omemoManager = omemoManager;
100        this.messageKey = aesKey;
101        this.initializationVector = iv;
102    }
103
104    /**
105     * Create a new OmemoMessageBuilder with random IV and AES key.
106     *
107     * @param omemoManager  omemoManager of our device.
108     * @param omemoStore    omemoStore.
109     * @param message       Messages body.
110     * @throws NoSuchPaddingException
111     * @throws BadPaddingException
112     * @throws InvalidKeyException
113     * @throws NoSuchAlgorithmException
114     * @throws IllegalBlockSizeException
115     * @throws UnsupportedEncodingException
116     * @throws NoSuchProviderException
117     * @throws InvalidAlgorithmParameterException
118     */
119    public OmemoMessageBuilder(OmemoManager omemoManager,
120                               OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> omemoStore, String message)
121            throws NoSuchPaddingException, BadPaddingException, InvalidKeyException, NoSuchAlgorithmException, IllegalBlockSizeException,
122            UnsupportedEncodingException, NoSuchProviderException, InvalidAlgorithmParameterException {
123        this.omemoManager = omemoManager;
124        this.omemoStore = omemoStore;
125        this.setMessage(message);
126    }
127
128    /**
129     * Create an AES messageKey and use it to encrypt the message.
130     * Optionally append the Auth Tag of the encrypted message to the messageKey afterwards.
131     *
132     * @param message content of the message
133     * @throws NoSuchPaddingException               When no Cipher could be instantiated.
134     * @throws NoSuchAlgorithmException             when no Cipher could be instantiated.
135     * @throws NoSuchProviderException              when BouncyCastle could not be found.
136     * @throws InvalidAlgorithmParameterException   when the Cipher could not be initialized
137     * @throws InvalidKeyException                  when the generated key is invalid
138     * @throws UnsupportedEncodingException         when UTF8 is unavailable
139     * @throws BadPaddingException                  when cipher.doFinal gets wrong padding
140     * @throws IllegalBlockSizeException            when cipher.doFinal gets wrong Block size.
141     */
142    public void setMessage(String message) throws NoSuchPaddingException, NoSuchAlgorithmException, NoSuchProviderException, InvalidAlgorithmParameterException, InvalidKeyException, UnsupportedEncodingException, BadPaddingException, IllegalBlockSizeException {
143        if (message == null) {
144            return;
145        }
146
147        // Encrypt message body
148        SecretKey secretKey = new SecretKeySpec(messageKey, KEYTYPE);
149        IvParameterSpec ivSpec = new IvParameterSpec(initializationVector);
150        Cipher cipher = Cipher.getInstance(CIPHERMODE, PROVIDER);
151        cipher.init(Cipher.ENCRYPT_MODE, secretKey, ivSpec);
152
153        byte[] body;
154        byte[] ciphertext;
155
156        body = (message.getBytes(StringUtils.UTF8));
157        ciphertext = cipher.doFinal(body);
158
159        byte[] clearKeyWithAuthTag = new byte[messageKey.length + 16];
160        byte[] cipherTextWithoutAuthTag = new byte[ciphertext.length - 16];
161
162        System.arraycopy(messageKey, 0, clearKeyWithAuthTag, 0, 16);
163        System.arraycopy(ciphertext, 0, cipherTextWithoutAuthTag, 0, cipherTextWithoutAuthTag.length);
164        System.arraycopy(ciphertext, ciphertext.length - 16, clearKeyWithAuthTag, 16, 16);
165
166        ciphertextMessage = cipherTextWithoutAuthTag;
167        messageKey = clearKeyWithAuthTag;
168    }
169
170    /**
171     * Add a new recipient device to the message.
172     *
173     * @param device recipient device
174     * @throws CryptoFailedException                when encrypting the messageKey fails
175     * @throws UndecidedOmemoIdentityException
176     * @throws CorruptedOmemoKeyException
177     */
178    public void addRecipient(OmemoDevice device) throws CryptoFailedException, UndecidedOmemoIdentityException, CorruptedOmemoKeyException {
179        addRecipient(device, false);
180    }
181
182    /**
183     * Add a new recipient device to the message.
184     * @param device recipient device
185     * @param ignoreTrust ignore current trust state? Useful for keyTransportMessages that are sent to repair a session
186     * @throws CryptoFailedException
187     * @throws UndecidedOmemoIdentityException
188     * @throws CorruptedOmemoKeyException
189     */
190    public void addRecipient(OmemoDevice device, boolean ignoreTrust) throws
191            CryptoFailedException, UndecidedOmemoIdentityException, CorruptedOmemoKeyException {
192        OmemoSession<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> session =
193                omemoStore.getOmemoSessionOf(omemoManager, device);
194
195        if (session != null) {
196            if (!ignoreTrust && !omemoStore.isDecidedOmemoIdentity(omemoManager, device, session.getIdentityKey())) {
197                // Warn user of undecided device
198                throw new UndecidedOmemoIdentityException(device);
199            }
200
201            if (!ignoreTrust && omemoStore.isTrustedOmemoIdentity(omemoManager, device, session.getIdentityKey())) {
202                // Encrypt key and save to header
203                CiphertextTuple encryptedKey = session.encryptMessageKey(messageKey);
204                keys.add(new OmemoVAxolotlElement.OmemoHeader.Key(encryptedKey.getCiphertext(), device.getDeviceId(), encryptedKey.isPreKeyMessage()));
205            }
206        }
207    }
208
209    /**
210     * Assemble an OmemoMessageElement from the current state of the builder.
211     *
212     * @return OmemoMessageElement
213     */
214    public OmemoVAxolotlElement finish() {
215        OmemoVAxolotlElement.OmemoHeader header = new OmemoVAxolotlElement.OmemoHeader(
216                omemoManager.getDeviceId(),
217                keys,
218                initializationVector
219        );
220        return new OmemoVAxolotlElement(header, ciphertextMessage);
221    }
222
223    /**
224     * Generate a new AES key used to encrypt the message.
225     *
226     * @return new AES key
227     * @throws NoSuchAlgorithmException
228     */
229    public static byte[] generateKey() throws NoSuchAlgorithmException {
230        KeyGenerator generator = KeyGenerator.getInstance(KEYTYPE);
231        generator.init(KEYLENGTH);
232        return generator.generateKey().getEncoded();
233    }
234
235    /**
236     * Generate a 16 byte initialization vector for AES encryption.
237     *
238     * @return iv
239     */
240    public static byte[] generateIv() {
241        SecureRandom random = new SecureRandom();
242        byte[] iv = new byte[16];
243        random.nextBytes(iv);
244        return iv;
245    }
246
247    public byte[] getCiphertextMessage() {
248        return ciphertextMessage;
249    }
250
251    public byte[] getMessageKey() {
252        return messageKey;
253    }
254}