package com.zeus.protection;
/**
 * ZXTM Java Extension for dynamically adding IP addresses to the list
 * of banned IP addresses in Service Protection classes.
 *
 * When the extension is invoked, it will append the IP address from
 * which the client connected to the list of banned addresses of one
 * Service protection class associated with the virtual server
 * that handled the connection.
 *
 * TrafficScript usage: java.run("com.zeus.protection.BlockAttack");
 *
 * @author: Michael Granzow, mgranzow@zeus.com
 */

import java.io.IOException;
import java.security.KeyStore;
import java.security.Provider;
import java.security.Security;
import java.security.cert.X509Certificate;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;

import javax.net.ssl.ManagerFactoryParameters;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactorySpi;
import javax.net.ssl.X509TrustManager;
import javax.servlet.GenericServlet;
import javax.servlet.ServletConfig;
import javax.servlet.ServletException;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;

import com.zeus.ZXTMServlet.ZXTMServletRequest;
import com.zeus.soap.zxtm._1_0.CatalogProtectionLocator;
import com.zeus.soap.zxtm._1_0.CatalogProtectionPort;
import com.zeus.soap.zxtm._1_0.VirtualServerLocator;
import com.zeus.soap.zxtm._1_0.VirtualServerPort;

public class BlockAttack extends GenericServlet {
 
   private static final long serialVersionUID = 1L;
   //! IPs that have already been banned are stored here.
   private static Set<String> bannedIPs =
      Collections.synchronizedSet(new HashSet<String>());
   //! The name of the admin user
   private static final String aname = "admin";
   //! The password of the admin user (please adapt this to your config)
   private static final String apw = "AdminneR";
   //! The object over which SOAP calls for Protection classes will be made
   private static CatalogProtectionPort cpport = null;
   //! The object over which SOAP calls for VServer classes will be made
   private static VirtualServerPort vsport = null;

   /**
    * Initialize the Servlet: Set up the trust manager and establish
    * the SOAP connections to the admin server.
    */
   public void init(ServletConfig cfg) throws ServletException {
      try {
         super.init(cfg);
         // Install the all-trusting trust manager:
         Security.addProvider(new MyProvider());
         Security.setProperty("ssl.TrustManagerFactory.algorithm",
         "TrustAllCertificates");
         // Set up the connections for the SOAP calls:
         String adminurl = "https://"+aname+":"+apw+"@127.0.0.1:9090/soap";
         CatalogProtectionLocator cplocator = new CatalogProtectionLocator();
         cplocator.setCatalogProtectionPortEndpointAddress(adminurl);
         cpport = cplocator.getCatalogProtectionPort();
         VirtualServerLocator vslocator = new VirtualServerLocator();
         vslocator.setVirtualServerPortEndpointAddress(adminurl);
         vsport = vslocator.getVirtualServerPort();
      } catch( Exception e ) {
         log("Exception in init: " + e.getMessage());
         throw new ServletException(e);
      }
   }

   /**
    * The service method: we treat all HTTP requests the
    * same way: block their origin.
    */
   public void service( ServletRequest req, ServletResponse res )
   throws ServletException, IOException
   {
      try {
         ZXTMServletRequest zreq = (ZXTMServletRequest)req;
         String srcip = (String)zreq.getAttribute("srcip");
         String vsname = (String)zreq.getAttribute("virtualserver");

         if( null == srcip || null == vsname ) {
            return; // ZXTM closed this connection anyway
         }
         if( !bannedIPs.add(srcip) ) {
            return;  // nothing to do for us
         }
         String[] vsnames = {vsname};
         String[] pnames = null;
         synchronized (vsport) {
            pnames = vsport.getProtection(vsnames);
         }
         if( pnames.length > 0 && pnames[0].length() > 0 ) {
            String[] protection = {pnames[0]};
            String[][] values = {{srcip}};
            synchronized (cpport) {
               cpport.addBannedAddresses(
                     protection,
                     values
                     );
            }
         } else {
            // log(vsname + " has no protection classes");
         }
      } catch( Exception e ) {
         log("Exception in service: " + e.getMessage());
      }
   }

   /* The following code disables certificate checking.
    * Use the Security.addProvider and Security.setProperty
    * calls to enable it */
   public static class MyProvider extends Provider {
      private static final long serialVersionUID = 1L;
      public MyProvider() {
         super( "MyProvider", 1.0, "Trust certificates" );
         put( "TrustManagerFactory.TrustAllCertificates",
               MyTrustManagerFactory.class.getName() );
      }
      protected static class MyTrustManagerFactory
      extends TrustManagerFactorySpi {
         public MyTrustManagerFactory() {}
         protected void engineInit( KeyStore keystore ) {}
         protected void engineInit(
               ManagerFactoryParameters mgrparams ) {}
         protected TrustManager[] engineGetTrustManagers() {
            return new TrustManager[] {
                  new MyX509TrustManager()
            };
         }
      }
      protected static class MyX509TrustManager implements X509TrustManager {
         public void checkClientTrusted(
               X509Certificate[] chain, String authType) {}
         public void checkServerTrusted(
               X509Certificate[] chain, String authType) {}
         public X509Certificate[] getAcceptedIssuers() {
            return null;
         }
      }
   }
}