/*
 * Copyright 2018 Scytl Secure Electronic Voting SA
 *
 * All rights reserved
 *
 * See our extended copyright notice in *file 'Copyright.txt' which is part of this source code package
 */
package com.scytl.products.ov.ag.infrastructure.filter;

import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.UnsupportedEncodingException;
import java.net.URLDecoder;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.function.Function;

import javax.servlet.ServletInputStream;
import javax.servlet.http.Cookie;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import javax.ws.rs.core.MediaType;

import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.StringUtils;
import org.jsoup.Jsoup;
import org.jsoup.safety.Whitelist;
import org.owasp.esapi.ESAPI;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.json.JsonSanitizer;

/**
 * A wrapper for the sanitized data of http request to avoid Cross-Site Scripting XSS attacks.
 */
public class SanitizerDataHttpServletRequestWrapper extends HttpServletRequestWrapper {

	private static final String CONTENT_TYPE = "Content-Type";

	private static final String TEXT_CSV = "text/csv";

	private static final String CSV_SEPARATOR = ",";

	private static final String QUERY_STRING = "QUERY STRING_";

	private static final String HEADER = "HEADER_";

	private static final String PARAMETER = "PARAMETER_";

	private static final String COOKIE = "COOKIE_";

	private static final String CSV = "CSV";

	private static final int LOG_MAX_STRING_WIDTH = 1024;
	
	private static final String SEPARATOR_LINE = "--------------------------------------------------------";

	private static final Logger LOG = LoggerFactory.getLogger(SanitizerDataHttpServletRequestWrapper.class);

	private final HttpServletRequest req;

	/**
	 * Constructor.
	 * 
	 * @param request the http request.
	 */
	public SanitizerDataHttpServletRequestWrapper(HttpServletRequest request) {
		super(request);
		req = request;
	}

	/**
	 * Strips any potential XSS threats out of the value.
	 * 
	 * @param value the concrete string value to be evaluated.
	 * @return a sanitized data.
	 */
	private static String stripXSS(String value) {
		if (value == null) {
			return null;
		}

		LOG.debug("Orginal value: " + value);

		String sanitizedData = value;

		// Use the ESAPI library to avoid encoded attacks.
		sanitizedData = ESAPI.encoder().canonicalize(sanitizedData);

		// Avoid null characters
		sanitizedData = StringUtils.replace(sanitizedData, "\0", "");

		// Clean out HTML
		sanitizedData = Jsoup.clean(sanitizedData, Whitelist.none());

		LOG.debug("New value: " + sanitizedData);
		LOG.debug(SEPARATOR_LINE);

		return sanitizedData;
		
	}

	// process headers
	@Override
	public String getHeader(String name) {
		LOG.debug(SEPARATOR_LINE);
		LOG.debug("Header " + name);
		String header = req.getHeader(name);
		String sanitizedHeader = stripXSS(header);
		checkIfDataIsSanitized(HEADER + name, header, sanitizedHeader);
		return header;
	}

	// logging
	private static void checkIfDataIsSanitized(String where, String originalValue, String sanitizedValue) {
		if (originalValue != null && sanitizedValue != null && !originalValue.isEmpty()
			&& !originalValue.equals(sanitizedValue)) {
			LOG.error("String was sanitized at {}.", where);
			LOG.error("Original value was: \"{}\".", StringUtils.abbreviate(originalValue, LOG_MAX_STRING_WIDTH));
			LOG.error("Sanitized value is: \"{}\".", StringUtils.abbreviate(sanitizedValue, LOG_MAX_STRING_WIDTH));
			throw new IllegalArgumentException("Data is not valid at " + where);
		}
	}

	// process input parameters in post
	@Override
	public ServletInputStream getInputStream() throws IOException {
		LOG.debug(SEPARATOR_LINE);
		LOG.debug("Input stream " + req.getContentType());
		try {
			String header = req.getHeader(CONTENT_TYPE);
			ServletInputStream inputStream = req.getInputStream();
			if (header != null) {
				String inputString;
				switch (header) {
				case MediaType.APPLICATION_JSON:
					inputString = IOUtils.toString(inputStream, StandardCharsets.UTF_8);
					String sanitizedJSON = JsonSanitizer.sanitize(inputString);
					checkIfDataIsSanitized(header, inputString, sanitizedJSON);
					return convertToServletInputStream(inputString);
				case TEXT_CSV:
					checkIfCSVIsSanitized(inputStream);
					return getInputStream(inputStream);
				default:
					inputString = IOUtils.toString(inputStream, StandardCharsets.UTF_8);
					String sanitizedIS = stripXSS(inputString);
					checkIfDataIsSanitized(header, inputString, sanitizedIS);
					return convertToServletInputStream(inputString);
				}
			}
			return inputStream;
		} catch (IOException e) {
			LOG.error("Error converting input stream to string in sanitizer data filter...", e);
			throw e;
		}
	}

	// input stream reading
	private ServletInputStream getInputStream(ServletInputStream inputStream) {
		ServletInputStream servletInputStream = new ServletInputStream() {
			@Override
			public int read() throws IOException {
				return inputStream.read();
			}
		};
		return servletInputStream;
	}

	// csv processing
	private void checkIfCSVIsSanitized(ServletInputStream inputStream) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(inputStream));
		br.lines().map(stripCSV);
	}

	// strip XSS from csv
	private static Function<String, String> stripCSV = line -> {
		String[] csv = line.split(CSV_SEPARATOR);
		for (int i = 0; i < csv.length; i++) {
			String originalValue = csv[i];
			csv[i] = stripXSS(originalValue);
			checkIfDataIsSanitized(CSV, originalValue, csv[i]);
		}
		return StringUtils.join(csv, CSV_SEPARATOR);
	};

	// converts string to stream
	public ServletInputStream convertToServletInputStream(String inputString) {
		ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(inputString.getBytes(StandardCharsets.UTF_8));
		ServletInputStream servletInputStream = new ServletInputStream() {
			@Override
			public int read() throws IOException {
				return byteArrayInputStream.read();
			}
		};
		return servletInputStream;
	}

	// process query string
	@Override
	public String getQueryString() {
		LOG.debug(SEPARATOR_LINE);
		LOG.debug("Query string");
		String queryString = req.getQueryString();
		Map<String, String[]> map;
		try {
			map = getQueryMap(queryString);
			checkIfMapValuesAreSanitized(QUERY_STRING, map);
		} catch (UnsupportedEncodingException e) {
			LOG.error("Error getting the query string map. Something is bad with the encoding.");
			throw new IllegalArgumentException("Data is not valid at " + QUERY_STRING, e);
		}
		return queryString;
	}

	// check query string
	private Map<String, String[]> getQueryMap(String queryString) throws UnsupportedEncodingException {
		Map<String, String[]> queryMap = new HashMap<>();
		if (queryString != null) {
			String[] pairs = queryString.split("&");
			for (String pair : pairs) {
				int idx = pair.indexOf("=");
				queryMap.put(URLDecoder.decode(pair.substring(0, idx), "UTF-8"),
					new String[] {URLDecoder.decode(pair.substring(idx + 1), "UTF-8") });
			}
		}
		return queryMap;
	}

	// process parameters
	@Override
	public String getParameter(String name) {
		LOG.debug(SEPARATOR_LINE);
		LOG.debug("Parameter " + name);
		String parameter = req.getParameter(name);
		String sanitizedParameter = stripXSS(parameter);
		checkIfDataIsSanitized(PARAMETER + parameter, parameter, sanitizedParameter);
		return parameter;
	}

	// process cookies
	@Override
	public Cookie[] getCookies() {
		LOG.debug(SEPARATOR_LINE);
		LOG.debug("Cookies");

		Cookie[] cookies = req.getCookies();
		for (Cookie cookie : cookies) {
			String cookieValue = cookie.getValue();
			String sanitizedCookie = stripXSS(cookieValue);
			cookie.setValue(sanitizedCookie);
			checkIfDataIsSanitized(COOKIE + cookie.getName(), cookieValue, sanitizedCookie);
		}
		return cookies;
	}

	// process parameter values
	@Override
	public String[] getParameterValues(String name) {
		LOG.debug(SEPARATOR_LINE);
		LOG.debug("Parameter value " + name);
		String[] values = req.getParameterValues(name);
		checkIfValuesAreSanitized(PARAMETER + name, values);
		return values;
	}

	// process parameter values
	@Override
	public Map<String, String[]> getParameterMap() {
		LOG.debug(SEPARATOR_LINE);
		LOG.debug("Parameter map");
		Map<String, String[]> map = req.getParameterMap();
		checkIfMapValuesAreSanitized(PARAMETER, map);
		return map;
	}

	// process map values
	private void checkIfMapValuesAreSanitized(String name, Map<String, String[]> map) {
		for (Entry<String, String[]> entry : map.entrySet()) {
			String key = entry.getKey();
			LOG.debug(name + " " + key);
			String[] values = map.get(key);
			checkIfValuesAreSanitized(name + key, values);
		}
	}

	// process array of values
	private String[] checkIfValuesAreSanitized(String name, String[] values) {
		String[] newValues = new String[values.length];
		for (int i = 0; i < values.length; i++) {
			newValues[i] = stripXSS(values[i]);
			checkIfDataIsSanitized(name, values[i], newValues[i]);
		}
		return newValues;
	}

}
