package org.thatmadhacker.finlaychat;

import java.io.File;
import java.io.FileWriter;
import java.io.PrintWriter;
import java.net.InetAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.util.ArrayList;
import java.util.List;
import java.util.Scanner;

import javax.crypto.SecretKey;

import com.thatmadhacker.utils.crypto.algo.ASymetric;
import com.thatmadhacker.utils.crypto.algo.Symetric;
import com.thatmadhacker.utils.misc.BASE64;

public class FinlayChat {

    public static final File USER_FILE = new File("user.profile");

    public static final int PORT = 7183;

    public static List<InetAddress> clients = new ArrayList<InetAddress>();

    public static List<String> messages = new ArrayList<String>();

    @SuppressWarnings("resource")
    public static void main(String[] args) throws Exception{

        Runtime.getRuntime().addShutdownHook(new Thread(new Runnable() {
            @Override
            public void run() {
                try {
                    UPnP.removePortMapping(PORT);
                }catch(Exception e) {
                    e.printStackTrace();
                }
            }
        }));
        UPnP.mapPort(PORT);
        new Thread(new Runnable() {
            @Override
            public void run() {
                try {
                    ServerSocket ss = new ServerSocket(PORT);
                    while(!ss.isClosed()) {
                        try {
                            Socket s = ss.accept();
                            Scanner scan = new Scanner(s.getInputStream());
                            if(!scan.nextLine().equals("UPDATE")) {
                                messages.add(scan.nextLine());
                            }else {
                                PrintWriter out = new PrintWriter(s.getOutputStream(),true);
                                for(String message : messages) {
                                    out.println(message);
                                }
                                out.close();
                            }
                            scan.close();
                        }catch(Exception e) {

                        }
                    }
                    ss.close();
                }catch(Exception e) {
                    e.printStackTrace();
                    System.exit(1);
                }
            }
        }).start();
        new Thread(new Runnable() {

            @Override
            public void run() {
                while(true) {
                    try {
                        updateMessages();

                        Thread.sleep(5000);
                    }catch(Exception e) {
                        e.printStackTrace();
                    }
                }
            }

        }).start();
        Scanner in = new Scanner(System.in);
        KeyPair pair;
        if(USER_FILE.exists()) {
            Scanner s = new Scanner(USER_FILE);
            String priv = s.nextLine();
            s.close();
            System.out.println("Enter decryption password!");
            String password = in.nextLine();
            SecretKey key = Symetric.genKey(password,"test123", 128,"AES");
            System.out.println("Attempting to decrypt private key!");
            String decryptedPriv = "";
            try {
                decryptedPriv = Symetric.decrypt(priv, key, "AES");
                System.out.println("Successfully decrypted private key!");
            }catch(Exception e) {
                System.out.println("WRONG PASSWORD!");
                System.exit(1);
            }
            PrivateKey privKey = ASymetric.getPrivateKeyFromByteArray(BASE64.decode(decryptedPriv.replaceAll("&l", System.lineSeparator())), "RSA");
            PublicKey pubKey = ASymetric.getPublicKeyFromRSAPrivateKey(privKey);
            pair = new KeyPair(pubKey,privKey);



        }else {
            USER_FILE.createNewFile();
            System.out.println("Enter a passcode to encrypt your key with!");
            String password = in.nextLine();
            SecretKey key = Symetric.genKey(password,"test123", 128,"AES");

            pair = ASymetric.genKeys("RSA",2048);

            String privKey = BASE64.encode(pair.getPrivate().getEncoded()).replaceAll(System.lineSeparator(), "&l");

            privKey = Symetric.encrypt(privKey, key, "AES");

            PrintWriter out = new PrintWriter(new FileWriter(USER_FILE));
            out.print(privKey);
            out.close();
        }

        String username = BASE64.encode(pair.getPublic().getEncoded()).replaceAll(System.lineSeparator(), "&l");

        updateMessages();


    }

    public static void updateMessages() throws Exception {
        for(InetAddress i : clients) {
            Socket s = new Socket(i,PORT);
            PrintWriter out = new PrintWriter(s.getOutputStream(),true);
            out.println("UPDATE");
            messages.clear();
            Scanner in = new Scanner(s.getInputStream());
            String m = in.nextLine();
            while(!m.equals("{END}")) {
                messages.add(m);
                m = in.nextLine();
            }
            in.close();
            s.close();
        }
    }

    public static void sendMessage(String to, String from, KeyPair pair, String message) throws Exception{
        PublicKey key = ASymetric.getPublicKeyFromByteArray(BASE64.decode(to.replaceAll("&l", System.lineSeparator())),"RSA");

        String encryptedMessage = ASymetric.encrypt(message, key, "RSA");

        String finalMessage = "";
        finalMessage += "from:"+from+"&to:"+to;
        String signedMessage = ASymetric.sign(encryptedMessage,pair.getPrivate(),"RSA");
        finalMessage += "&message:"+signedMessage;

        transmit(finalMessage);

    }

    private static void transmit(String message) throws Exception{
        for(InetAddress i : clients) {
            Socket s = new Socket(i,PORT);
            PrintWriter out = new PrintWriter(s.getOutputStream(),true);
            out.println("PUT");
            out.println(message);
            out.close();
            s.close();
        }
    }

}