package com.frontend.controllers;
import com.CryptoAlgorithms.DiffieHellman.DHAgreementModel;
import com.CryptoAlgorithms.DiffieHellman.DHAlgorithm;
import com.exceptions.CreatingDHAgreementException;
import com.frontend.models.request.DiffieHellman.FinalizeDHRequest;
import com.frontend.models.response.DefaultResponse;
import com.frontend.models.response.DiffieHellman.DHAgreementResponse;
import com.frontend.models.response.ErrorResponse;
import com.frontend.models.response.Response;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestMethod;
import org.springframework.web.bind.annotation.ResponseBody;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
@Controller
@RequestMapping(value = "/keyExchange")
public class KeyExchangeController {
private final DHAlgorithm dhAlgorithm;
private Map<UUID, DHAgreementModel> keyExchangeSessions = new HashMap<UUID, DHAgreementModel>();
private Map<UUID, byte[]> secretKeys = new HashMap<UUID, byte[]>();
@Autowired
public KeyExchangeController(DHAlgorithm dhAlgorithm) {
this.dhAlgorithm = dhAlgorithm;
}
@RequestMapping(value = "/initialize", method = RequestMethod.GET)
@ResponseBody
public Response initializeExchange() {
try {
DHAgreementModel dhAgreementModel = dhAlgorithm.generateInitialNumbers();
UUID keyExchangeSessionId = UUID.randomUUID();
keyExchangeSessions.put(keyExchangeSessionId, dhAgreementModel);
return new DHAgreementResponse(dhAgreementModel.getP(), dhAgreementModel.getG(), dhAgreementModel.getPublicA(), keyExchangeSessionId);
} catch (CreatingDHAgreementException e) {
return new ErrorResponse(e.getMessage());
}
}
@RequestMapping(value = "/finalize", method = RequestMethod.GET)
@ResponseBody
public Response finalizeExchange(@RequestBody(required = true) FinalizeDHRequest finalizeDHRequest) {
if (finalizeDHRequest.getKeyExchangeSessionId() == null || !keyExchangeSessions.containsKey(finalizeDHRequest.getKeyExchangeSessionId())) {
return new ErrorResponse("Invalid keyExchangeSessionId.");
}
DHAgreementModel dhAgreementModel = keyExchangeSessions.get(finalizeDHRequest.getKeyExchangeSessionId());
try {
byte[] secretKey = dhAlgorithm.generateSecret(dhAgreementModel.get_a(), finalizeDHRequest.getPublicB());
secretKeys.put(finalizeDHRequest.getKeyExchangeSessionId(), secretKey);
keyExchangeSessions.remove(finalizeDHRequest.getKeyExchangeSessionId());
return new DefaultResponse();
} catch (CreatingDHAgreementException e) {
return new ErrorResponse(e.getMessage());
}
}
}