package sample;

import javax.security.auth.callback.Callback;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.callback.NameCallback;
import javax.security.auth.callback.PasswordCallback;
import javax.security.auth.callback.UnsupportedCallbackException;
import javax.security.sasl.AuthorizeCallback;
import javax.security.sasl.Sasl;
import javax.security.sasl.SaslClient;
import javax.security.sasl.SaslException;
import javax.security.sasl.SaslServer;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

public class App {

    private SaslServer saslServer;

    public static void main(String[] args) throws SaslException {
        new App().init().start();
    }

    private App init() {
        try {
            saslServer = Sasl.createSaslServer("CRAM-MD5", "my_server", "local", null, new ServerHandler());
        } catch (SaslException e) {
            e.printStackTrace();
        }
        return this;
    }

    private static class ClientHandler implements CallbackHandler {

        private String username;
        private char[] password;


        public void setUsername(String username) {
            this.username = username;
        }

        public void setPassword(char[] password) {
            this.password = password;
        }

        public void handle(Callback[] cbs) throws IOException, UnsupportedCallbackException {
            for (Callback cb : cbs) {
                if (cb instanceof NameCallback) {

                    System.out.println("Client - NameCallback");

                    NameCallback nc = (NameCallback) cb;
                    nc.setName(username);
                } else if (cb instanceof PasswordCallback) {

                    System.out.println("Client - PasswordCallback");

                    PasswordCallback pc = (PasswordCallback) cb;
                    pc.setPassword(password);
                }
            }
        }
    }

    private static class ServerHandler implements CallbackHandler {

        private Map<String,String> authorizedUsers = new HashMap<String,String>();

        {
            authorizedUsers.put("energy", "mypass");
            authorizedUsers.put("user", "userpass");
        }


        public void handle(Callback[] cbs) throws IOException, UnsupportedCallbackException {
            String user = null;
            for (Callback cb : cbs) {
                if (cb instanceof AuthorizeCallback) {

                    System.out.println("Server - AuthorizeCallback");

                    AuthorizeCallback ac = (AuthorizeCallback) cb;
                    ac.setAuthorized(true);

                } else if (cb instanceof NameCallback) {

                    System.out.println("Server - NameCallback");

                    NameCallback nc = (NameCallback) cb;
                    if (authorizedUsers.containsKey(nc.getDefaultName())) {
                        user = nc.getDefaultName();
                        nc.setName(nc.getDefaultName());
                    }
                } else if (cb instanceof PasswordCallback) {
                    System.out.println("Server - PasswordCallback");

                    PasswordCallback pc = (PasswordCallback) cb;
                    if (user != null) {
                        pc.setPassword(authorizedUsers.get(user).toCharArray());
                    }
                }
            }
        }
    }

    private void start() throws SaslException {

        byte[] challenge;
        byte[] response;

        ClientHandler clientHandler = new ClientHandler();

        clientHandler.setUsername("energy");
        clientHandler.setPassword("mypass".toCharArray());
        SaslClient sc = Sasl.createSaslClient(new String[]{"CRAM-MD5"}, null, "my_server", "local", null, clientHandler);

        challenge = saslServer.evaluateResponse(new byte[0]);
        response = sc.evaluateChallenge(challenge);

        try {
            saslServer.evaluateResponse(response);
            if (saslServer.isComplete()) {

                System.out.println("Authentication successful.");
            }
        } catch (SaslException e) {
            e.printStackTrace();
        }
    }
}