// Copyright (c) Trend Micro Inc. 2016, All Rights Reserved

package com.trendmicro.ds.restapisamples;

import java.util.Arrays;
import java.util.Calendar;
import java.util.Comparator;
import java.util.GregorianCalendar;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.logging.Level;
import java.util.logging.Logger;

import com.trendmicro.ds.platform.rest.object.monitoring.TenantHostProtectionElement;
import com.trendmicro.ds.platform.rest.object.monitoring.TenantHostProtectionListing;
import com.trendmicro.ds.platform.rest.object.monitoring.TenantModuleUsageElement;
import com.trendmicro.ds.restapisamples.DeepSecurityClient;

/**
 * Billing
 *
 * This sample application shows how to use the Deep Security Manager REST API
 * to retrieve computer protection information.
 */
public class Billing {

	private static final String USAGE = "Billing <DSM URL> <username> <password>\n" +
			"This sample application shows how to use the Deep Security Manager REST API\n" +
			"to retrieve computer protection information.\n";

	private static final Logger logger = Logger.getLogger(Billing.class.getName());

	/** Mapping from known protection module codes to a more human-friendly string */
	private static final Map<String, String> moduleNameToFriendlyName;

	static {
		moduleNameToFriendlyName = new HashMap<>();
		moduleNameToFriendlyName.put("WRS", "Web Reputation");
		moduleNameToFriendlyName.put("IM", "Integrity Monitoring");
		moduleNameToFriendlyName.put("LI", "Log Inspection");
		moduleNameToFriendlyName.put("DPI", "Intrusion Prevention");
		moduleNameToFriendlyName.put("FW", "Firewall");
		moduleNameToFriendlyName.put("AM", "Anti-Malware");
	}

	private static final double MILLISECONDS_PER_HOUR = 60.0 * 60.0 * 1000.0;

	/**
	 * The main entry point.
	 *
	 * @param args see USAGE.
	 */
	public static void main(String[] args) {
		if (args.length < 3) {
			logger.warning(USAGE);
			return;
		}

		new Billing().run(args[0], args[1], args[2]);
	}

	/**
	 * Run the billing sample.
	 *
	 * @param baseUrl base URL for the Deep Security Manager REST API (for example: {@code https://dsm.example.com/rest})
	 * @param username username to use when accessing the Deep Security Manager REST API
	 * @param password password to use when accessing the Deep Security Manager REST API
	 */
	private void run(String baseUrl, String username, String password) {
		try (DeepSecurityClient client = new DeepSecurityClient(baseUrl, username, password, null)) {

			/*
			 * You should <strong>only</strong> use this in a secure development environment, and <strong>never</strong>
			 * use it over an untrusted network. Disabling the trust manager will turn off the validation process that
			 * checks whether the server you are connecting to is the one you expect, and can result in your connection
			 * being hijacked, your credentials stolen, and general catastrophe. Use with great caution!
			 */
			// client.disableTrustManager();

			// Set up search date range to search for the previous day
			Calendar start = new GregorianCalendar();
			start.set(Calendar.MILLISECOND, 0);
			start.set(Calendar.SECOND, 0);
			start.set(Calendar.MINUTE, 0);
			start.set(Calendar.HOUR, 0);
			start.add(Calendar.DAY_OF_YEAR, -1);

			Calendar end = (Calendar)start.clone();
			end.add(Calendar.DAY_OF_YEAR, 1);

			// Get all the protection records for the last day for the tenant
			// with ID 0 - the primary tenant.
			TenantHostProtectionListing protectionListing = client.listHostProtection(null, Integer.valueOf(0), start.getTime(), end.getTime());

			// Calculate and display information about protection received
			calculateProtectionReceived(protectionListing.getProtectionElements(), start, end);

		} catch (DeepSecurityClient.ClientException e) {
			logger.log(Level.SEVERE, "There was an error communicating with the server: " + e.getMessage());
		} catch (Exception e) {
			// Some other error happened, usually related to network communication problems.
			logger.log(Level.SEVERE, "There was an error communicating with the server.", e);
		}
	}

	/**
	 * Calculate the amount of protection time received by protection module.
	 *
	 * @param protectionList
	 * 		the list of protection records retrieved from Deep Security Manager
	 *
	 * @param periodStart
	 * 		the start of the time interval over which to calculate protection received
	 *
	 * @param periodEnd
	 * 		the end of the time interval over which to calculate protection received
	 */
	@SuppressWarnings("boxing")
	private void calculateProtectionReceived(List<TenantHostProtectionElement> protectionList, Calendar periodStart, Calendar periodEnd) {
		long minTime = periodStart.getTimeInMillis();
		long maxTime = periodEnd.getTimeInMillis();

		// Total protection received. This is tracked separately as it's possible
		// the agent was active on a computer with no protection modules active
		long totalProtectionReceived = 0;

		// Total protection received broken down by module
		Map<String, Long> protectionReceivedByModule = new HashMap<>();

		for (TenantHostProtectionElement protectionRecord : protectionList) {
			// Protection start time is the most recent of the time in the protection
			// record and the start of the billing time interval
			long protectionStart = Math.max(minTime, protectionRecord.getProtectionStartDate().getTime());

			// Protection stop time is either the value from the protection record
			// (if there is one it means protection stopped during the time interval)
			// or the end of the billing time interval
			long protectionStop;
			if (protectionRecord.getProtectionStopDate() == null) {
				protectionStop = maxTime;
			} else {
				protectionStop = protectionRecord.getProtectionStopDate().getTime();
			}

			// Increment total protection duration
			long protectionDuration = protectionStop - protectionStart;
			totalProtectionReceived += protectionDuration;

			// Increment protection duration for each protection module
			for (TenantModuleUsageElement moduleUsage : protectionRecord.getModuleUsageList()) {
				if (moduleUsage.isModuleEnabled()) {
					// Only need to do the calculation if the module was actually enabled
					Long currentModuleTotal = protectionReceivedByModule.get(moduleUsage.getModuleName());
					if (currentModuleTotal == null) {
						currentModuleTotal = protectionDuration;
					} else {
						currentModuleTotal = currentModuleTotal + protectionDuration;
					}
					protectionReceivedByModule.put(moduleUsage.getModuleName(), currentModuleTotal);
				}
			}

		}

		// Display the calculated overall and per-module protection times in hours
		// A real billing system would feed the numbers computed here in to some cost calculation.
		logger.info(String.format("Billing time period from %s to %s\n", periodStart.getTime().toString(), periodEnd.getTime().toString()));
		logger.info(String.format("Total protection duration: %,.2f hours\n", totalProtectionReceived / MILLISECONDS_PER_HOUR));
		for (Map.Entry<String, Long> mapEntry : protectionReceivedByModule.entrySet()) {
			// Try to map internal module name to a friendly string
			String friendlyName = moduleNameToFriendlyName.get(mapEntry.getKey());
			if (friendlyName == null) {
				friendlyName = mapEntry.getKey();
			}

			logger.info(String.format("%,.2f hours of %s protection received\n", mapEntry.getValue().doubleValue() / MILLISECONDS_PER_HOUR, friendlyName));
		}

		// Some billing scenarios bill by peak number of protected computers in the billing period
		int peakProtectedComputers = computePeakProtectedComputers(protectionList);
		logger.info(String.format("Maximum number of computers protected at the same time: %d", peakProtectedComputers));

		/*
		 This is a fairly simple demo to show what is possible. More advanced scenarios include:
		 1. Rounding up protection time for each computer to the nearest hour.

		    One thing to be careful of when rounding up is that changing the
		    protection a computer receives causes the old protection record
		    to receive a protection end date and a new protection record
		    is created with:
		    	-start date equal to the end date in the old record
		    	-the new protection module information
		    	-no protection end date
		    This may or may not need to be accounted for depending on the
		    rounding requirements.

		 2. Breaking down protection based on cloud instance type.

		 	Deep Security Manager must be configured with a cloud connector
		 	for this information to be available in TenantModuleUsageElement.

		 */
	}

	/**
	 * Compute the maximum number of computers that were protected simultaneously.
	 * This requires computing the maximum number of overlapping protection time intervals.
	 * The idea is as follows: <br/>
	 * - create a flat list of every protection start or stop event <br/>
	 * - sort the list from earliest to latest <br/>
	 * - traverse the list, incrementing the protected computer count for
	 * 		each protection start event and decrementing it for each
	 * 		protection stop event <br/>
	 * - the desired value is the maximum protected computer count
	 */
	private int computePeakProtectedComputers(List<TenantHostProtectionElement> protectionList) {
		long now = System.currentTimeMillis();

		// First, create a flat list of everything that happened (start or stop) during the time period.
		ProtectionEvent[] protectionEvents = new ProtectionEvent[protectionList.size() * 2];
		for (int i = 0; i < protectionList.size(); ++i) {
			TenantHostProtectionElement protectionRecord = protectionList.get(i);
			protectionEvents[2*i] = new ProtectionEvent(protectionRecord.getProtectionStartDate().getTime(), true);
			if (protectionRecord.getProtectionStopDate() == null) {
				protectionEvents[2*i + 1] = new ProtectionEvent(now, false);
			} else {
				protectionEvents[2*i + 1] = new ProtectionEvent(protectionRecord.getProtectionStopDate().getTime(), false);
			}
		}

		// Now sort the list so we have the order in which the actions happened
		Arrays.sort(protectionEvents, new Comparator<ProtectionEvent>() {
			/* Comparator to sort ProtectionEvents from earliest to latest. Sorting is defined so that if
			 * one interval ends at the same time another begins, the end of one is ordered before the start
			 * of the other. */
			@Override
			public int compare(ProtectionEvent o1, ProtectionEvent o2) {
				long timeDiff = o1.time - o2.time;
				if (timeDiff > 0) {
					return 1;
				} else if (timeDiff < 0) {
					return -1;
				} else if (o1.isStart == o2.isStart) {
					return 0;
				} else if (o1.isStart && !o2.isStart) {
					return 1;
				} else if (!o1.isStart && o2.isStart) {
					return -1;
				}
				return 0;
			}
		});

		// Now it's simple. Any time an protection record starts, one extra computer came online, so
		// add 1 to the total. When a record ends, a computer is no longer being protected, so
		// subtract 1 from the total.
		int peakProtectedComputers = 0;
		int currentPeak = 0;
		for (int i = 0; i < protectionEvents.length; ++i) {
			if (protectionEvents[i].isStart) {
				++currentPeak;
				peakProtectedComputers = Math.max(currentPeak, peakProtectedComputers);
			} else {
				--currentPeak;
			}
			assert currentPeak >= 0;
		}

		assert currentPeak == 0;

		return peakProtectedComputers;
	}

	/** Simple class to track the beginning or the end of a protection record. */
	private static class ProtectionEvent {
		final long time;
		final boolean isStart;

		public ProtectionEvent(long time, boolean isStart) {
			this.time = time;
			this.isStart = isStart;
		}

	}

}