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

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.ArrayList;
import java.util.Arrays;
import java.util.Base64;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.function.Consumer;

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.Label;
import org.expeditee.encryption.items.surrogates.Surrogate;
import org.expeditee.encryption.items.surrogates.Label.LabelResult;
import org.expeditee.gui.Frame;
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 String SURROGATE_TERMINATOR = 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.resolveKey(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
        public void outputFrame(Frame frame) throws IOException {
                if (_writer == null) { return; }
                
                preOutputFrame();
                writeHeader(frame);
                
                // write item
                writeItemData(frame);
                writeTerminator();
                
                // write lines and constraints
                writeLineData();
                writeTerminator();
                writeConstraintData();
                writeTerminator();
                
                writeLine(SessionStats.getFrameEventList(frame));
        }
        
        @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
        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);
        }
        
        @Override
        protected void writeLineData() throws IOException {
                super.writeLineData();
                
                writeSurrogateLineData();
        }
        
        @Override
        protected void writeClass(Item toWrite) throws IOException {
                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");
                
                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) {
                EncryptionDetail encryptionDetail = toWrite.getEncryptionDetailForTag(tag + "");
                switch (encryptionDetail) {
                case UnencryptedOnSave:
                        super.writeTag(toWrite, param, tags, tag);
                        break;
                case ReencryptOnSave:
                        Method toRun = tags.get(tag);
                        Class<?> declarer = toRun.getDeclaringClass();
                        LabelResult res = Label.resolveKey(toWrite.getEncryptionLabel());
                        if (declarer.isAssignableFrom(toWrite.getClass()) && res == LabelResult.SuccessResolveLabelToKey) {
                                SecretKey key = new SecretKeySpec(res.key, SymmetricAlgorithm);
                                try {
                                        Object o = toRun.invoke(toWrite, new Object[] {});
                                        o = Conversion.ConvertToExpeditee(toRun, o);
                                        if (o != null) {
                                                if (o instanceof List) {
                                                        for (Object line: (List<?>)o) {
                                                                byte[] encryptedBytes = EncryptSymmetric(line.toString().getBytes(), key);
                                                                writeLine(tag + "E", Base64.getEncoder().encodeToString(encryptedBytes));
                                                        }
                                                } else {
                                                        byte[] encryptedBytes = EncryptSymmetric(o.toString().getBytes(), key);
                                                        writeLine(tag + "E", Base64.getEncoder().encodeToString(encryptedBytes));
                                                }
                                        }
                                } catch (IllegalAccessException | IllegalArgumentException | InvocationTargetException e) {
                                        e.printStackTrace();
                                } catch (IOException e) {
                                        e.printStackTrace();
                                }
                        }
                        break;
                case UseUndecipheredValueOnSave:
                        try {
                                writeLine(tag + "E", encryptionDetail.getUndecipheredValue());
                        } catch (IOException e) {
                                e.printStackTrace();
                        }
                        break;
                }
        }
        
        private void writeItemData(Frame frame) throws IOException {
                // write each item in the frame
                for (Item i : frame.getItemsToSave()) {
                        assert (!(i instanceof Line));
                        writeItem(i);
                }
                
                for (final Item i: frame.getBodyItemsWithInsufficientPermissions()) {
                        assert (!(i instanceof Line));
                        writeItem(i);
                }
                
                writeSurrogateItemsData(frame);
        }
        
        private void writeSurrogateItemsData(Frame frame) throws IOException {
                List<Surrogate> surrogates = new ArrayList<Surrogate>();
                
                Consumer<? super Item> store = item -> surrogates.addAll(item.getSurrogates());
                frame.getItemsToSave().forEach(store);
                frame.getBodyItemsWithInsufficientPermissions().forEach(store);
                
                if (!surrogates.isEmpty()) {
                        // If there is some surrogates, write the surrogate 
                        // terminator, then write out the surrogates.
                        writeSurrogateTerminator();
                        
                        for (Surrogate surrogate: surrogates) {
                                String output = surrogate.toString();
                                writeLine(output);
                        }
                }
        }

        private void writeSurrogateLineData() throws IOException {
                List<Surrogate> surrogateLines = new ArrayList<Surrogate>();
                
                for(Item lineEnd: _lineEnds) {
                        List<Line> lines = lineEnd.getLines();
                        for (Line line: lines) {
                                surrogateLines.addAll(line.getSurrogates());
                        }
                }
                
                if (!surrogateLines.isEmpty()) {
                        // If there is some surrogates, write the surrogate 
                        // terminator, then write out the surrogates.
                        writeSurrogateTerminator();
                        
                        for (Surrogate surrogate: surrogateLines) {
                                String output = surrogate.toString();
                                writeLine(output);
                        }
                }
        }
        
        private void writeSurrogateTerminator() throws IOException {
                writeLine(SURROGATE_TERMINATOR + nl);
        }
        
        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;
                }
        }       
}