source: trunk/src/org/expeditee/encryption/io/EncryptedExpWriter.java@ 1415

package org.expeditee.encryption.io;

import java.io.IOException;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.util.Arrays;
import java.util.Base64;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.stream.Collectors;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

import org.expeditee.encryption.CryptographyConstants;
import org.expeditee.encryption.items.surrogates.EncryptionDetail;
import org.expeditee.encryption.items.surrogates.EncryptionDetail.Type;
import org.expeditee.encryption.items.surrogates.Label;
import org.expeditee.encryption.items.surrogates.Label.LabelResult;
import org.expeditee.gui.Frame;
import org.expeditee.gui.Frame.BodyType;
import org.expeditee.gui.MessageBay;
import org.expeditee.io.Conversion;
import org.expeditee.io.ExpWriter;
import org.expeditee.items.Item;
import org.expeditee.items.Line;
import org.expeditee.items.Text;
import org.expeditee.settings.identity.secrets.KeyList;
import org.expeditee.stats.SessionStats;

public class EncryptedExpWriter extends ExpWriter implements CryptographyConstants {
        private SecretKey key;
        
        private final String label;
        
        private static final String ENCRYPTED_EXP_FLAG = "EncryptedExp";
        private static final String nl = "\n";
        private static final String labelProfile = "Profile";
        private static final String labelNone = "None";
        protected static final char TERMINATOR = 'Z';
        protected static final String TERMINATOR_WITH_CONTINUATION = TERMINATOR + "...";

        public EncryptedExpWriter(String encryptionLabel) {
                label = encryptionLabel;
                if (label.equals(labelNone)) {
                        return;
                }
                
                if (label.equals(labelProfile)) {
                        // obtain personal key
                        Text text = KeyList.PersonalKey.get();
                        List<String> data = text.getData();
                        if (data != null && !data.isEmpty()) {
                                byte[] keyBytes = Base64.getDecoder().decode(data.get(0));
                                key = new SecretKeySpec(keyBytes, SymmetricAlgorithm);
                        }
                } else {
                        LabelResult res = Label.getLabel(label);
                        if (res == LabelResult.SuccessResolveLabelToKey) {
                                byte[] keyBytes = res.key;
                                key = new SecretKeySpec(keyBytes, SymmetricAlgorithm);
                        } else if (res == LabelResult.ErrorUnableToFindLabel) {
                                MessageBay.errorMessage(res.toString() + encryptionLabel);
                                key = null;
                        } else {
                                MessageBay.errorMessage(res.toString());
                        }
                }
        }
        
        @Override
        protected void preOutputFrame() {
                try {
                        String line = ENCRYPTED_EXP_FLAG + label + nl;
                        _writer.write(line);
                        _stringWriter.append(line);
                } catch (final IOException e) { 
                        e.printStackTrace();
                }
        }
        
        @Override
        public void outputFrame(Frame frame) throws IOException {
                if (_writer == null) { return; }
                
                preOutputFrame();
                writeHeader(frame);
                
                // write item
                writeItemData(frame, true);
                writeTerminator();
                // write lines and constraints
                writeLineData();
                writeTerminator();
                writeConstraintData();
                
                // write surrogate items
                if (frame.hasSurrogates()) {
                        writeSurrogateTerminator();
                        writeItemData(frame, false);
                        writeTerminator();
                } else {
                        writeTerminator();
                }
                                
                writeLine(SessionStats.getFrameEventList(frame));
        }
        
        @Override
        protected void writeLine(String line) throws IOException {
                // do not write empty lines
                if (line == null) { return; }
                
                String toWrite;
                if (key == null && label.equals(labelNone)) {
                        toWrite = line + nl;
                } else {
                        // prepare line to write out
                        byte[] encrypted = EncryptSymmetric(line.getBytes(), key);
                        toWrite = Base64.getEncoder().encodeToString(encrypted) + nl;
                }
                                
                // output
                _writer.write(toWrite);
                _stringWriter.append(toWrite);
        }
        
        protected void writeClass(Item toWrite) throws IOException {
                if (toWrite instanceof Text && toWrite.getText().equals("Content for Label Two")) {
                        System.err.println("EncryptedExpWriter::writeClass::breakpoint");
                } else if (toWrite instanceof Text && toWrite.getText().equals("Encrypted") && toWrite.getPrimary().getText().equals("Content for Label One")) {
                        System.err.println("EncryptedExpWriter::writeClass::breakpoint");
                }
                
                LinkedHashMap<Character,Method> itemTags = new LinkedHashMap<Character, Method>(getItemTags());
                LinkedHashMap<String,Method> itemTagsExt = new LinkedHashMap<String, Method>(getItemTagsExt());
                
                writeTag(toWrite, new Object[] {}, itemTags, 'S');
                writeTag(toWrite, new Object[] {}, itemTagsExt, "_el");
                
                if (toWrite.isSurrogate()) {
                        writeLine("SurrogateFor " + toWrite.getPrimary().getID());
                }
                
                itemTags.remove('S');
                itemTagsExt.remove("_el");
                                
                writeTags(toWrite, new Object[] {}, itemTags);
                writeTags(toWrite, new Object[] {}, itemTagsExt);
        }
        
        @Override
        protected <T> void writeTag(Item toWrite, Object[] param, LinkedHashMap<T, Method> tags, T tag) {
                if (toWrite.isSurrogate() && toWrite.isTagInherited(tag + "")) {
                        return;
                }
                EncryptionDetail encryptionDetail = toWrite.getEncryptionDetailForTag(tag + "");
                Type encryptionDetailType = encryptionDetail.getEncryptionDetailType();
                                
                switch(encryptionDetailType) {
                case UnencryptedOnSave:
                        writeTagUnencryptedOnSave(toWrite, param, tags, tag);
                        break;
                case InheritanceCheckOnSave:
                        writeTagInheritanceCheckOnSave(toWrite, tags, tag);
                        break;
                case ReencryptOnSave:
                        writeTagReencryptOnSave(toWrite, tags, tag);
                        break;
                case UseUndecipheredValueOnSave:
                        writeTagUseUndecipheredValueOnSave(toWrite, tags, tag);
                        break;
                }
        }
        
        private void writeItemData(Frame frame, boolean primaryItems) throws IOException {
                // write each item in the frame
                BodyType bodyType = primaryItems ? BodyType.PrimaryBody : BodyType.SurrogateBody;
                for (Item i : frame.getItemsToSave(bodyType)) {
                        assert (!(i instanceof Line));
                        writeItem(i);
                }
                
                for (Item i: frame.getBodyItemsWithInsufficientPermissions()) {
                        assert (!(i instanceof Line));
                        writeItem(i);
                }
        }
        
        private <T> void writeTagInheritanceCheckOnSave(Item toWrite, LinkedHashMap<T, Method> tags, T tag) {   
                List<Item> surrogateItems = toWrite.getSurrogates().stream().collect(Collectors.toList());
                Function<Item, Boolean> isTagInherited = surrogate -> surrogate.isTagInherited(tag + "");
                BinaryOperator<Boolean> trueExists = (a, b) -> a || b;
                boolean surrogatesInherit = surrogateItems.stream().map(isTagInherited).collect(Collectors.reducing(trueExists)).get();
                boolean userHasKey = Label.getLabel(toWrite.getEncryptionLabel()) == LabelResult.SuccessResolveLabelToKey;
                
                if (!surrogatesInherit && userHasKey) {
                        EncryptionDetail reencryptOnSave = new EncryptionDetail(EncryptionDetail.Type.ReencryptOnSave);
                        toWrite.setEncryptionDetailForTag(tag + "", reencryptOnSave);
                        writeTagReencryptOnSave(toWrite, tags, tag);
                } else {
                        EncryptionDetail unencryptedOnSave = new EncryptionDetail(EncryptionDetail.Type.UnencryptedOnSave);
                        toWrite.setEncryptionDetailForTag(tag + "", unencryptedOnSave);
                        writeTagUnencryptedOnSave(toWrite, new Object[] {}, tags, tag);
                }
        }
        
        private <T> void writeTagReencryptOnSave(Item toWrite, LinkedHashMap<T, Method> tags, T tag) {
                Method toRun = tags.get(tag);
                Class<?> declarer = toRun.getDeclaringClass();
                LabelResult res = Label.getLabel(toWrite.getEncryptionLabel());
                if (declarer.isAssignableFrom(toWrite.getClass()) && res == LabelResult.SuccessResolveLabelToKey) {
                        try {
                                Object o = toRun.invoke(toWrite, new Object[] {});
                                o = Conversion.ConvertToExpeditee(toRun, o);
                                if (o == null) {
                                        return;
                                }
                                if (o instanceof List) {
                                        for (Object line: (List<?>) o) {
                                                writeLineEnc(tag, line, res.key);
                                        }
                                } else {
                                        writeLineEnc(tag, o, res.key);
                                }
                        } catch (IllegalAccessException | IllegalArgumentException | InvocationTargetException e) {
                                e.printStackTrace();
                        } catch (IOException e) {
                                e.printStackTrace();
                        }
                }
        }
        
        private <T> void writeTagUseUndecipheredValueOnSave(Item toWrite, LinkedHashMap<T, Method> tags, T tag) {
                EncryptionDetail encryptionDetail = toWrite.getEncryptionDetailForTag(tag + "");
                try {
                        writeLine(tag + "E", encryptionDetail.getUndecipheredValue());
                } catch (IOException e) {
                        e.printStackTrace();
                }
        }
        
        private <T> void writeTagUnencryptedOnSave(Item toWrite, Object[] param, LinkedHashMap<T, Method> tags, T tag) {
                super.writeTag(toWrite, param, tags, tag);
        }
        
        private void writeSurrogateTerminator() throws IOException {
                writeLine(TERMINATOR_WITH_CONTINUATION + nl);
        }
        
        private <T> void writeLineEnc(T tag, Object line, byte[] keyBytes) throws IOException {
                SecretKey key = new SecretKeySpec(keyBytes, SymmetricAlgorithm);
                byte[] lineEncryptedBytes = EncryptSymmetric(line.toString().getBytes(), key);
                line = Base64.getEncoder().encodeToString(lineEncryptedBytes);
                writeLine(tag + "E", line.toString());
        }
        
        private static byte[] EncryptSymmetric(byte[] toEncrypt, SecretKey key) {
                try {
                        Cipher cipher = Cipher.getInstance(SymmetricAlgorithm + SymmetricAlgorithmParameters);
                        cipher.init(Cipher.ENCRYPT_MODE, key);
                        //could use modulus
                        int length = (int) ((Math.ceil(toEncrypt.length / 16f)) * 16);
                        byte[] toEncryptSizeAdjusted = Arrays.copyOf(toEncrypt, length);
                        byte[] result = cipher.doFinal(toEncryptSizeAdjusted);
                        return result;
                } catch (final NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException
                                | IllegalBlockSizeException | BadPaddingException e) {
                        e.printStackTrace();
                        return null;
                }
        }
}