/*
 * Copyright (c) 2010, Isode Limited, London, England.
 * All rights reserved.
 */
/*
 * Copyright (c) 2010, Remko Tronçon.
 * All rights reserved.
 */
package com.isode.stroke.sasl;

import com.isode.stroke.base.ByteArray;
import com.isode.stroke.stringcodecs.Base64;
import com.isode.stroke.stringcodecs.HMACSHA1;
import com.isode.stroke.stringcodecs.PBKDF2;
import com.isode.stroke.stringcodecs.SHA1;
import java.text.Normalizer;
import java.text.Normalizer.Form;
import java.util.HashMap;
import java.util.Map;

public class SCRAMSHA1ClientAuthenticator extends ClientAuthenticator {

    static String escape(String s) {
        String result = "";
        for (int i = 0; i < s.length(); ++i) {
            if (s.charAt(i) == ',') {
                result += "=2C";
            } else if (s.charAt(i) == '=') {
                result += "=3D";
            } else {
                result += s.charAt(i);
            }
        }
        return result;
    }

    public SCRAMSHA1ClientAuthenticator(String nonce) {
        this(nonce, false);
    }
    public SCRAMSHA1ClientAuthenticator(String nonce, boolean useChannelBinding) {
        super(useChannelBinding ? "SCRAM-SHA-1-PLUS" : "SCRAM-SHA-1");
        step = Step.Initial;
        clientnonce = nonce;
        this.useChannelBinding = useChannelBinding;
    }

    public void setTLSChannelBindingData(ByteArray channelBindingData) {
        tlsChannelBindingData = channelBindingData;
    }

    public ByteArray getResponse() {
        if (step.equals(Step.Initial)) {
            return ByteArray.plus(getGS2Header(), getInitialBareClientMessage());
        } else if (step.equals(Step.Proof)) {
            ByteArray clientKey = HMACSHA1.getResult(saltedPassword, new ByteArray("Client Key"));
            ByteArray storedKey = SHA1.getHash(clientKey);
            ByteArray clientSignature = HMACSHA1.getResult(storedKey, authMessage);
            ByteArray clientProof = clientKey;
            byte[] clientProofData = clientProof.getData();
            for (int i = 0; i < clientProofData.length; ++i) {
                clientProofData[i] ^= clientSignature.getData()[i];
            }
            ByteArray result = getFinalMessageWithoutProof().append(",p=").append(Base64.encode(clientProof));
            return result;
        } else {
            return null;
        }
    }

    public boolean setChallenge(ByteArray challenge) {
        if (step.equals(Step.Initial)) {
            if (challenge == null) {
                return false;
            }
            initialServerMessage = challenge;

            Map<Character, String> keys = parseMap(initialServerMessage.toString());

            // Extract the salt
            ByteArray salt = Base64.decode(keys.get('s'));

            // Extract the server nonce
            String clientServerNonce = keys.get('r');
            if (clientServerNonce.length() <= clientnonce.length()) {
                return false;
            }
            String receivedClientNonce = clientServerNonce.substring(0, clientnonce.length());
            if (!receivedClientNonce.equals(clientnonce)) {
                return false;
            }
            serverNonce = new ByteArray(clientServerNonce.substring(clientnonce.length()));


            // Extract the number of iterations
            int iterations = 0;
            try {
                iterations = Integer.parseInt(keys.get('i'));
            } catch (NumberFormatException e) {
                return false;
            }
            if (iterations <= 0) {
                return false;
            }

            ByteArray channelBindData = new ByteArray();
            if (useChannelBinding && tlsChannelBindingData != null) {
                channelBindData = tlsChannelBindingData;
            }

            // Compute all the values needed for the server signature
            saltedPassword = PBKDF2.encode(new ByteArray(SASLPrep(getPassword())), salt, iterations);
            authMessage = getInitialBareClientMessage().append(",").append(initialServerMessage).append(",").append(getFinalMessageWithoutProof());
            ByteArray serverKey = HMACSHA1.getResult(saltedPassword, new ByteArray("Server Key"));
            serverSignature = HMACSHA1.getResult(serverKey, authMessage);

            step = Step.Proof;
            return true;
        } else if (step.equals(step.Proof)) {
            ByteArray result = new ByteArray("v=").append(new ByteArray(Base64.encode(serverSignature)));
            step = Step.Final;
            return challenge != null && challenge.equals(result);
        } else {
            return true;
        }
    }

    private String SASLPrep(String source) {
        return Normalizer.normalize(source, Form.NFKC); /* FIXME: Implement real SASLPrep */
    }

    private Map<Character, String> parseMap(String s) {
        HashMap<Character, String> result = new HashMap<Character, String>();
        if (s.length() > 0) {
            char key = '~'; /* initialise so it'll compile */
            String value = "";
            int i = 0;
            boolean expectKey = true;
            while (i < s.length()) {
                if (expectKey) {
                    key = s.charAt(i);
                    expectKey = false;
                    i++;
                } else if (s.charAt(i) == ',') {
                    result.put(key, value);
                    value = "";
                    expectKey = true;
                } else {
                    value += s.charAt(i);
                }
                i++;
            }
            result.put(key, value);
        }
        return result;
    }

    private ByteArray getInitialBareClientMessage() {
        String authenticationID = SASLPrep(getAuthenticationID());
        return new ByteArray("n=" + escape(authenticationID) + ",r=" + clientnonce);
    }

    private ByteArray getGS2Header() {

        ByteArray channelBindingHeader = new ByteArray("n");
	if (tlsChannelBindingData != null) {
		if (useChannelBinding) {
			channelBindingHeader = new ByteArray("p=tls-unique");
		}
		else {
			channelBindingHeader = new ByteArray("y");
		}
	}
	return new ByteArray().append(channelBindingHeader).append(",").append(getAuthorizationID().isEmpty() ? new ByteArray() : new ByteArray("a=" + escape(getAuthorizationID()))).append(",");
    }

    private ByteArray getFinalMessageWithoutProof() {
        ByteArray channelBindData = new ByteArray();
	if (useChannelBinding && tlsChannelBindingData != null) {
		channelBindData = tlsChannelBindingData;
	}
	return new ByteArray("c=" + Base64.encode(new ByteArray(getGS2Header()).append(channelBindData)) + ",r=" + clientnonce).append(serverNonce);
    }

    private enum Step {

        Initial,
        Proof,
        Final
    };
    private Step step;
    private String clientnonce = "";
    private ByteArray initialServerMessage = new ByteArray();
    private ByteArray serverNonce = new ByteArray();
    private ByteArray authMessage = new ByteArray();
    private ByteArray saltedPassword = new ByteArray();
    private ByteArray serverSignature = new ByteArray();
    private boolean useChannelBinding;
    private ByteArray tlsChannelBindingData;
}