package org.owasp.webgoat;

import java.io.IOException;
import java.io.PrintWriter;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.Enumeration;
import java.util.List;
import java.util.Locale;
import java.util.TimeZone;

import javax.servlet.ServletContext;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpSession;

import org.owasp.webgoat.lessons.AbstractLesson;
import org.owasp.webgoat.lessons.WelcomeScreen;
import org.owasp.webgoat.lessons.admin.WelcomeAdminScreen;
import org.owasp.webgoat.session.Course;
import org.owasp.webgoat.session.ErrorScreen;
import org.owasp.webgoat.session.Screen;
import org.owasp.webgoat.session.UserTracker;
import org.owasp.webgoat.session.WebSession;

/**
 *  Copyright (c) 2002 Free Software Foundation developed under the custody of the Open Web
 *  Application Security Project (http://www.owasp.org) This software package is published by OWASP
 *  under the GPL. You should read and accept the LICENSE before you use, modify and/or redistribute
 *  this software.
 *
 * @author     Jeff Williams <a href="http://www.aspectsecurity.com">Aspect Security</a>
 * @created    October 28, 2003
 */
public class HammerHead extends HttpServlet
{
	/**
	 *  Description of the Field
	 */
	protected static SimpleDateFormat httpDateFormat;

	/**
	 *  Description of the Field
	 */
	protected WebSession mySession;
	
	/**
	 * Set the session timeout to be 2 days
	 */
	private final static int sessionTimeoutSeconds = 60 * 60 * 24 * 2;
	//private final static int sessionTimeoutSeconds = 1;

	
	/**
	 * Properties file path
	 */
	public static String propertiesPath = null;
	
	

	/**
	 *  Description of the Method
	 *
	 * @param  request               Description of the Parameter
	 * @param  response              Description of the Parameter
	 * @exception  IOException       Description of the Exception
	 * @exception  ServletException  Description of the Exception
	 */
	public void doGet( HttpServletRequest request, HttpServletResponse response ) throws IOException, ServletException
	{
		doPost( request, response );
	}


	/**
	 *  Description of the Method
	 *
	 * @param  request               Description of the Parameter
	 * @param  response              Description of the Parameter
	 * @exception  IOException       Description of the Exception
	 * @exception  ServletException  Description of the Exception
	 */
	public void doPost( HttpServletRequest request, HttpServletResponse response ) throws IOException, ServletException
	{
		Screen screen = null;

		try
		{
			//System.out.println( "HH Entering doPost: " );
			//System.out.println( "	- HH request "   + request);
			//System.out.println( "	- HH principle: "   + request.getUserPrincipal() );
			//setCacheHeaders(response, 0);
			ServletContext context = getServletContext();

			// FIXME: If a response is written by updateSession(), do not call makeScreen() and writeScreen()
			mySession = updateSession( request, response, context );
			if (response.isCommitted())
				return;
						
			// Note: For the lesson to track the status, we need to update the lesson tracker object
			//       from the screen.createContent() method.  The create content is the only point
			//		 where the lesson "knows" what has happened.  To track it at a latter point would
			// 		 require the lesson to have memory.
			screen = makeScreen( mySession );	// This calls the lesson's handleRequest()
			if (response.isCommitted())
				return;
			
			// if the screen parameter exists, the screen was visited via the menu categories,
			// we won't count these as visits. The user may be able to manipulate the counts
			// by specifying the screen parameter using a proxy.  Good for them!
			String fromMenus = mySession.getParser().getRawParameter( WebSession.SCREEN, null );
			if ( fromMenus == null )
			{
				// if the show source parameter exists, don't add the visit
				fromMenus = mySession.getParser().getRawParameter( WebSession.SHOW, null );
				if ( fromMenus == null  )
				{
					screen.getLessonTracker( mySession ).incrementNumVisits();
				}
			}

			// log the access to this screen for this user
			UserTracker userTracker = UserTracker.instance();
			userTracker.update( mySession, screen );
			log( request, screen.getClass().getName() + " | " + mySession.getParser().toString() );

			// Redirect the request to our View servlet
			String userAgent = request.getHeader("user-agent");
			String clientBrowser = "Not known!";
			if (userAgent != null)
			{
			    clientBrowser = userAgent;
			}
			request.setAttribute("client.browser", clientBrowser);
			request.getSession().setAttribute("websession", mySession);
			request.getSession().setAttribute("course", mySession.getCourse());
			
			request.getRequestDispatcher(getViewPage(mySession)).forward(request, response);
		}
		catch ( Throwable t )
		{
			t.printStackTrace();
			log( "ERROR: " + t );
			screen = new ErrorScreen( mySession, t );
		}
		finally
		{
			try
			{
				this.writeScreen( screen, response );
			}
			catch ( Throwable thr )
			{
				thr.printStackTrace();
				log( request, "Could not write error screen: " + thr.getMessage() );
			}
			//System.out.println( "HH Leaving doPost: " );
		}
	}

	private String getViewPage(WebSession webSession)
	{
		String page;
		
		// If this session has not seen the landing page yet, go there instead.
		HttpSession session = webSession.getRequest().getSession();
		if (session.getAttribute("welcomed") == null)
		{
			session.setAttribute("welcomed", "true");
			page = "/webgoat.jsp";
		}
		else
			page = "/main.jsp";
		
		return page;
	}
	
	/**
	 *  Description of the Method
	 *
	 * @param  session  Description of the Parameter
	 */
	private void dumpSession( HttpSession session )
	{
		Enumeration enumerator = session.getAttributeNames();

		while ( enumerator.hasMoreElements() )
		{
			String name = (String) enumerator.nextElement();
			Object value = session.getAttribute( name );
			System.out.println( "Name: " + name );
			System.out.println( "Value: " + value );
		}
	}


	/**
	 *  Description of the Method
	 *
	 * @param  date  Description of the Parameter
	 * @return       RFC 1123 http date format
	 */
	protected static String formatHttpDate( Date date )
	{
		synchronized ( httpDateFormat )
		{
			return httpDateFormat.format( date );
		}
	}


	/**
	 *  Return information about this servlet
	 *
	 * @return    The servletInfo value
	 */
	public String getServletInfo()
	{
		return "WebGoat is sponsored by Aspect Security.";
	}
	
	/**
	 *  Return properties path 
	 *  
	 *  @return servlet context path + WEB_INF
	 */
	public void init() throws ServletException
	{
		httpDateFormat = new SimpleDateFormat( "EEE, dd MMM yyyyy HH:mm:ss z", Locale.US );
		httpDateFormat.setTimeZone( TimeZone.getTimeZone( "GMT" ) );
		propertiesPath = getServletContext().getRealPath(
				"." + System.getProperty("file.separator")+ "WEB-INF" + "/webgoat.properties");
	}


	/**
	 *  Description of the Method
	 *
	 * @param  request  Description of the Parameter
	 * @param  message  Description of the Parameter
	 */
	public void log( HttpServletRequest request, String message )
	{
		String output = new Date() + " | " + request.getRemoteHost() + ":" + request.getRemoteAddr() + " | " + message;
		log( output );
		System.out.println( output );
	}

	public List getCategories()
	{
		Course course = mySession.getCourse();
		
		// May need to clone the List before returning it.
		//return new ArrayList(course.getCategories());
		return course.getCategories();
	}
	
/*
 	public List getLessons(Category category, String role)
	{
		Course course = mySession.getCourse();
		
		// May need to clone the List before returning it.
		//return new ArrayList(course.getLessons(category, role));
		return course.getLessons(category, role);
	}
*/
	

	/**
	 *  Description of the Method
	 *
	 * @param  s  Description of the Parameter
	 * @return    Description of the Return Value
	 */
	protected Screen makeScreen( WebSession s )
	{
		Screen screen = null;
		int scr = s.getCurrentScreen();
		Course course = s.getCourse();

		if ( s.isUser() || s.isChallenge() )
		{
			if ( scr == WebSession.WELCOME )
			{
				screen = new WelcomeScreen( s );
			}
			else
			{
				AbstractLesson lesson = course.getLesson( s, scr, AbstractLesson.USER_ROLE );
				if ( lesson == null && s.isHackedAdmin() )
				{
					// If admin was hacked, let the user see some of the admin screens
					lesson = course.getLesson( s, scr, AbstractLesson.HACKED_ADMIN_ROLE );
				}

				if ( lesson != null )
				{
					screen = lesson;
					
					// We need to do some bookkeeping for the hackable admin interface.
					// This is the only place we can tell if the user successfully hacked the hackable
					// admin and has actually accessed an admin screen.  You need BOTH pieces of information
					// in order to satisfy the remote admin lesson.
					
					s.setHasHackableAdmin( screen.getRole() );
					
					lesson.handleRequest( s );
					s.setCurrentMenu( lesson.getCategory().getRanking() );
				}
				else
				{
					screen = new ErrorScreen( s, "Invalid screen requested.  Try: http://localhost/WebGoat/attack" );
				}
			}
		}
		else if ( s.isAdmin() )
		{
			if ( scr == WebSession.WELCOME )
			{
				screen = new WelcomeAdminScreen( s );
			}
			else
			{
				// Admin can see all roles.
				// FIXME: should be able to pass a list of roles.
				AbstractLesson lesson = course.getLesson( s, scr, AbstractLesson.ADMIN_ROLE );
				if ( lesson == null )
				{
					lesson = course.getLesson( s, scr, AbstractLesson.HACKED_ADMIN_ROLE );
				}
				if ( lesson == null )
				{
					lesson = course.getLesson( s, scr, AbstractLesson.USER_ROLE );
				}

				if ( lesson != null )
				{
					screen = lesson;
					
					// We need to do some bookkeeping for the hackable admin interface.
					// This is the only place we can tell if the user successfully hacked the hackable
					// admin and has actually accessed an admin screen.  You need BOTH pieces of information
					// in order to satisfy the remote admin lesson.
					
					s.setHasHackableAdmin( screen.getRole() );

					lesson.handleRequest( s );
					s.setCurrentMenu( lesson.getCategory().getRanking() );
				}
				else
				{
					screen = new ErrorScreen( s, "Invalid screen requested.  Try Setting Admin to false or Try: http://localhost/WebGoat/attack" );
				}
			}
		}
		
		return ( screen );
	}


	/**
	 *  This method sets the required expiration headers in the response for a given RunData object.
	 *  This method attempts to set all relevant headers, both for HTTP 1.0 and HTTP 1.1.
	 *
	 * @param  response  The new cacheHeaders value
	 * @param  expiry    The new cacheHeaders value
	 */
	protected static void setCacheHeaders( HttpServletResponse response, int expiry )
	{
		if ( expiry == 0 )
		{
			response.setHeader( "Pragma", "no-cache" );
			response.setHeader( "Cache-Control", "no-cache" );
			response.setHeader( "Expires", formatHttpDate( new Date() ) );
		}
		else
		{
			Date expiryDate = new Date( System.currentTimeMillis() + expiry );
			response.setHeader( "Expires", formatHttpDate( expiryDate ) );
		}
	}


	/**
	 *  Description of the Method
	 *
	 * @param  request   Description of the Parameter
	 * @param  response  Description of the Parameter
	 * @param  context   Description of the Parameter
	 * @return           Description of the Return Value
	 */
	protected WebSession updateSession( HttpServletRequest request, HttpServletResponse response, ServletContext context )
			throws IOException
	{
		HttpSession hs;
		hs = request.getSession( true );

		//System.out.println( "HH Entering Session_id: " + hs.getId() );
		// dumpSession( hs );
		// Get our session object out of the HTTP session
		WebSession session = null;
		Object o = hs.getAttribute( WebSession.SESSION );

		if ( ( o != null ) && o instanceof WebSession )
		{
			session = (WebSession) o;
		}
		else
		{
			// Create new custom session and save it in the HTTP session
			//System.out.println( "HH Creating new WebSession: " );
			session = new WebSession( this, context );
			hs.setAttribute( WebSession.SESSION, session );
			// reset timeout
			hs.setMaxInactiveInterval( sessionTimeoutSeconds );

		}

		session.update( request, response, this.getServletName() );

		// to authenticate
		//System.out.println( "HH Leaving Session_id: " + hs.getId() );
		//dumpSession( hs );
		return ( session );
	}


	/**
	 *  Description of the Method
	 *
	 * @param  s                Description of the Parameter
	 * @param  response         Description of the Parameter
	 * @exception  IOException  Description of the Exception
	 */
	protected void writeScreen( Screen s, HttpServletResponse response ) throws IOException
	{
		response.setContentType( "text/html" );

		PrintWriter out = response.getWriter();

		if ( s == null )
		{
			s = new ErrorScreen( mySession, "Page to display was null" );
		}
		
		// set the content-length of the response.
		// Trying to avoid chunked-encoding.  (Aspect required)
		response.setContentLength( s.getContentLength() );
		response.setHeader("Content-Length",s.getContentLength()+"");

		s.output( out );
		out.close();
	}
}

