import java.io.*;
import javax.servlet.*;
import javax.servlet.http.*;

/**
 * Response Wrapper that captures output from the ServletOutputStream
 * and the PrinterWriter.  This allows interception and modification 
 * of a response's output in a Filter.
 *
 * @author Daniel J. Troesser
 */

public class OutputCaptureResponseWrapper 
        extends HttpServletResponseWrapper {

    private ByteArrayServletOutputStream basos;
    private PrintWriter pw;
    private String contentType;

    /**
     * Configure a new response wrapper.
     *
     * @param response the response we are wrapping
     */
    public OutputCaptureResponseWrapper(HttpServletResponse response) {
        super(response);
        basos = new ByteArrayServletOutputStream();
        pw = new PrintWriter(basos);
    }
    
    // Override setContentType() so we can store and retrieve the content type
    // later.  This might be useful for Filters that want to process only
    // one type of content.
    public void setContentType(String contentType) {
        this.contentType = contentType;
        super.setContentType(contentType);
    }
    
    public String getContentType() {
        return contentType;
    }
    
    /**
     * Returns the underlying output byte array as a String
     *
     * @return the underlying output byte array as a String
     */
    public String getOutputAsString() {
        return basos.toString();
    }

    /**
     * Returns the underlying output byte array
     *
     * @return the underlying output byte array
     */
    public byte[] getOutputAsByteArray() {
        return basos.getByteArrayOutputStream().toByteArray();
    }

    // Override the methods involving the response's Writer or
    // OutputStream.

    public void flushBuffer() throws IOException {
        pw.flush();
    }

    public int getBufferSize() {
        return basos.getSize();
    }

    public ServletOutputStream getOutputStream() {
        return basos;
    }

    public PrintWriter getWriter() throws IOException {
        return pw;
    }

    public void reset() {
        resetBuffer();
        super.reset();
    }

    public void resetBuffer() {
        basos.reset();
    }

    public void setBufferSize(int size) throws IllegalStateException {
        try {
            basos.setSize(size);
        } catch (IOException ioe) {
            throw new IllegalStateException(ioe.getMessage());
        }
    }
}