org.zaproxy.zap.extension.plugnhack.db.MessageTable.java Source code

Introduction

Here is the source code for org.zaproxy.zap.extension.plugnhack.db.MessageTable.java

Source

/*
 * Zed Attack Proxy (ZAP) and its related class files.
 *
 * ZAP is an HTTP/HTTPS proxy for assessing web application security.
 *
 * Copyright 2014 The ZAP Development Team
 *
 * 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.plugnhack.db;

import java.io.IOException;
import java.io.InputStream;
import java.io.StringWriter;
import java.nio.charset.StandardCharsets;
import java.sql.CallableStatement;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Timestamp;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import net.sf.json.JSONObject;
import net.sf.json.JSONSerializer;
import org.apache.commons.io.IOUtils;
import org.apache.log4j.Logger;
import org.hsqldb.jdbc.JDBCClob;
import org.parosproxy.paros.db.DatabaseException;
import org.parosproxy.paros.db.DbUtils;
import org.parosproxy.paros.db.paros.ParosAbstractTable;
import org.zaproxy.zap.extension.plugnhack.ClientMessage;
import org.zaproxy.zap.extension.plugnhack.ClientMessage.State;

public class MessageTable extends ParosAbstractTable {

    private static final String TABLE = "PLUGNHACK_MESSAGE";
    private static final Logger logger = Logger.getLogger(MessageTable.class);

    private PreparedStatement psInsert;
    private PreparedStatement psUpdate;
    private PreparedStatement psDelete;

    private PreparedStatement psGetAllData = null;

    private CallableStatement psGetIdLastInsert = null;

    @Override
    protected void reconnect(Connection conn) throws DatabaseException {
        try {
            if (!DbUtils.hasTable(conn, TABLE)) {
                // need to create the tables
                PreparedStatement stmt = conn.prepareStatement("CREATE CACHED TABLE " + TABLE + " ("
                        + "id BIGINT GENERATED BY DEFAULT AS IDENTITY (START WITH 1), "
                        + "timestamp TIMESTAMP NOT NULL," + "client_id VARCHAR(255) NOT NULL,"
                        + "state TINYINT NOT NULL," + "message CLOB(16M) NOT NULL," + "changed BOOLEAN NOT NULL"
                        + ")");
                DbUtils.executeAndClose(stmt);
            }

            psInsert = conn.prepareStatement("INSERT INTO " + TABLE
                    + " (timestamp, client_id, state, message, changed) " + "VALUES (?,?,?,?,?)");

            psUpdate = conn.prepareStatement(
                    "UPDATE " + TABLE + " SET message = ?, state = ?, changed = TRUE WHERE id = ?");

            psDelete = conn.prepareStatement("DELETE FROM " + TABLE + " WHERE id = ?");

            psGetAllData = conn.prepareStatement("SELECT * FROM " + TABLE);

            psGetIdLastInsert = conn.prepareCall("CALL IDENTITY();");
        } catch (SQLException e) {
            throw new DatabaseException(e);
        }
    }

    public synchronized void insert(ClientMessage cmsg) throws SQLException {
        psInsert.setTimestamp(1, new Timestamp(cmsg.getReceived().getTime()));
        psInsert.setString(2, cmsg.getClientId());
        psInsert.setInt(3, cmsg.getState().ordinal());
        psInsert.setClob(4, new JDBCClob(cmsg.getJson().toString()));
        psInsert.setBoolean(5, cmsg.isChanged());

        psInsert.executeUpdate();

        ResultSet rs = psGetIdLastInsert.executeQuery();
        rs.next();
        cmsg.setIndex(rs.getLong(1));
        rs.close();
    }

    public synchronized void update(ClientMessage cmsg) throws SQLException {
        psUpdate.setClob(1, new JDBCClob(cmsg.getJson().toString()));
        psUpdate.setInt(2, cmsg.getState().ordinal());
        psUpdate.setLong(3, cmsg.getIndex());

        psUpdate.executeUpdate();
    }

    public synchronized void delete(ClientMessage cmsg) throws SQLException {
        psDelete.setLong(1, cmsg.getIndex());
        psDelete.executeUpdate();
    }

    public synchronized List<ClientMessage> list() throws SQLException {
        ArrayList<ClientMessage> list = new ArrayList<ClientMessage>();

        ResultSet rs = psGetAllData.executeQuery();
        while (rs.next()) {
            try {
                InputStream in = rs.getClob("message").getAsciiStream();
                StringWriter w = new StringWriter();
                IOUtils.copy(in, w, StandardCharsets.UTF_8);
                ClientMessage cmsg = new ClientMessage(rs.getString("client_id"),
                        (JSONObject) JSONSerializer.toJSON(w.toString()));
                cmsg.setIndex(rs.getLong("id"));
                cmsg.setReceived(new Date(rs.getTimestamp("timestamp").getTime()));
                cmsg.setChanged(rs.getBoolean("changed"));
                cmsg.setState(State.values()[rs.getInt("state")]);
                list.add(cmsg);
            } catch (IOException e) {
                logger.error(e.getMessage(), e);
            }
        }
        rs.close();
        return list;
    }
}