StatsAPI.java

/*
 /*
 * Zed Attack Proxy (ZAP) and its related class files.
 * 
 * ZAP is an HTTP/HTTPS proxy for assessing web application security.
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); 
 * you may not use this file except in compliance with the License. 
 * You may obtain a copy of the License at 
 * 
 *   http://www.apache.org/licenses/LICENSE-2.0 
 *   
 * Unless required by applicable law or agreed to in writing, software 
 * distributed under the License is distributed on an "AS IS" BASIS, 
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
 * See the License for the specific language governing permissions and 
 * limitations under the License. 
 */
package org.zaproxy.zap.extension.stats;

import java.util.Map;
import java.util.Map.Entry;
import java.util.TreeMap;

import net.sf.json.JSON;
import net.sf.json.JSONArray;
import net.sf.json.JSONObject;

import org.apache.commons.httpclient.URI;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Text;
import org.zaproxy.zap.extension.api.ApiAction;
import org.zaproxy.zap.extension.api.ApiException;
import org.zaproxy.zap.extension.api.ApiImplementor;
import org.zaproxy.zap.extension.api.ApiResponse;
import org.zaproxy.zap.extension.api.ApiResponseElement;
import org.zaproxy.zap.extension.api.ApiResponseList;
import org.zaproxy.zap.extension.api.ApiResponseSet;
import org.zaproxy.zap.extension.api.ApiView;
import org.zaproxy.zap.model.SessionStructure;
import org.zaproxy.zap.utils.Stats;
import org.zaproxy.zap.utils.XMLStringUtil;

public class StatsAPI extends ApiImplementor {

	private static final String PREFIX = "stats";
	
	private static final String ACTION_CLEAR_STATS = "clearStats";
	
	private static final String VIEW_STATS = "stats";
	private static final String VIEW_SITE_STATS = "siteStats";
	private static final String VIEW_ALL_SITES_STATS = "allSitesStats";

	private static final String PARAM_KEY_PREFIX = "keyPrefix";
	private static final String PARAM_SITE = "site";
	
	private ExtensionStats extension;

	public StatsAPI(ExtensionStats extension) {
		this.extension = extension;
		this.addApiAction(new ApiAction(ACTION_CLEAR_STATS, null, new String[] {PARAM_KEY_PREFIX}));
		
		this.addApiView(new ApiView(VIEW_STATS, null, new String[] { PARAM_KEY_PREFIX }));
		this.addApiView(new ApiView(VIEW_ALL_SITES_STATS, null, new String[] {PARAM_KEY_PREFIX }));
		this.addApiView(new ApiView(VIEW_SITE_STATS, new String[] {PARAM_SITE}, new String[] {PARAM_KEY_PREFIX }));
	}

	@Override
	public String getPrefix() {
		return PREFIX;
	}

	@Override
	public ApiResponse handleApiAction(String name, JSONObject params)
			throws ApiException {
		if (ACTION_CLEAR_STATS.equals(name)) {
			Stats.clear(this.getParam(params, PARAM_KEY_PREFIX, ""));
			return ApiResponseElement.OK;
			
		} else {
			throw new ApiException(ApiException.Type.BAD_ACTION);
		}
	}
	
	@Override
	public ApiResponse handleApiView(String name, JSONObject params)
			throws ApiException {
		ApiResponse result = null;
		InMemoryStats memStats = extension.getInMemoryStats();
		if (memStats == null) {
			throw new ApiException(ApiException.Type.DOES_NOT_EXIST);
		}

		if (VIEW_STATS.equals(name)) {
			Map<String, String> map = new TreeMap<>();

			for (Entry<String, Long> stat : memStats.getStats(this.getParam(params, PARAM_KEY_PREFIX, "")).entrySet()) {
				map.put(stat.getKey(), stat.getValue().toString());
			}
			result = new ApiResponseSet<String>(name, map);

		} else if (VIEW_ALL_SITES_STATS.equals(name)) {
			result = new ApiResponseList(name);
			for (Entry<String, Map<String, Long>> stats :
				memStats.getAllSiteStats(this.getParam(params, PARAM_KEY_PREFIX, "")).entrySet()) {
				((ApiResponseList)result).addItem(new SiteStatsApiResponse(stats.getKey(), stats.getValue()));
			}

		} else if (VIEW_SITE_STATS.equals(name)) {
			String site = params.getString(PARAM_SITE);
			URI siteURI;
			
			try {
				siteURI = new URI(site, true);
				site = SessionStructure.getHostName(siteURI);
			} catch (Exception e) {
				throw new ApiException(ApiException.Type.ILLEGAL_PARAMETER, PARAM_SITE);
			}
			String scheme = siteURI.getScheme();
			if (scheme == null || (!scheme.equalsIgnoreCase("http") && !scheme.equalsIgnoreCase("https"))) {
				throw new ApiException(ApiException.Type.ILLEGAL_PARAMETER, PARAM_SITE);
			}

			result = new SiteStatsApiResponse(site, 
					memStats.getSiteStats(site, this.getParam(params, PARAM_KEY_PREFIX, "")));

		} else {
			throw new ApiException(ApiException.Type.BAD_VIEW);
		}
		return result;
	}
	
	private static class SiteStatsApiResponse extends ApiResponseList {

		private String site;
		private Map<String, Long> stats;
		
		public SiteStatsApiResponse(String site, Map<String, Long> stats) {
			super("statistics");
			this.site = site;
			this.stats = new TreeMap<>(stats);
			this.addItem(new ApiResponseSet<Long>(site, this.stats));
		}
		
		@Override
		public void toXML(Document doc, Element parent) {
			parent.setAttribute("type", "list");
			Element els = doc.createElement("site");
			Text texts = doc.createTextNode(XMLStringUtil.escapeControlChrs(this.site));
			els.appendChild(texts);
			parent.appendChild(els);
			
			for (Entry<String, Long> stat : this.stats.entrySet()) {
				Element el = doc.createElement("statistic");
				el.setAttribute("type", "set");
				
				Element elk = doc.createElement("key");
				Text textk = doc.createTextNode(XMLStringUtil.escapeControlChrs(stat.getKey()));
				elk.appendChild(textk);
				el.appendChild(elk);

				Element elv = doc.createElement("value");
				Text textv = doc.createTextNode(XMLStringUtil.escapeControlChrs(stat.getValue().toString()));
				elv.appendChild(textv);
				el.appendChild(elv);
				
				parent.appendChild(el);
			}
		}

		@Override
		public JSON toJSON() {
			JSONObject jo = new JSONObject();
			JSONArray array = new JSONArray();
			for (ApiResponse resp: this.getItems()) {
				if (resp instanceof ApiResponseElement) {
					array.add(((ApiResponseElement)resp).getValue());
				} else {
					array.add(resp.toJSON());
				}
			}
			// Use the site name instead of 'statistics'
			jo.put(this.site, array);
			return jo;
		}
	}
}