C:/work/RevealedSystems/EyeSpy-RS/eyespy/src/objects/storage/sqlite3Impl.py

Go to the documentation of this file.

# Copyright (C) 2008-2009  Mark Hanegraaff
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from objects.logic.common.exception import exception
from pysqlite2 import dbapi2 as sqlite
from objects.storage.storageInterface import storageInterface
from objects.logic.common.logger import logger
import time
import shutil
from objects.logic.common.globalStrings import globalStrings

import base64
import os



## The sqlite3Impl class is the implementation of the storageInterface using SQLLite. 
# SQLLite is a public domain SQL database available from (www.sqlite.org) 
class sqlite3Impl(storageInterface):
    
    def __init__(self, storageFile):
        
        # error call back function, should take 1 string parameter
        # representing the actual error
        self.errorCallBackFunc = None        
        
        self.cur  = None
        self.conn = None
        
        self.recID = 0
        
        self.storageFile = storageFile
        self.successful = False
        
            
            
    ## Creates the SQLite storage file. Note that an empty database file (named capture.tpl) must be created ahead of time. 
    # The capture file is created using the following SQLite statement:
    #
    # BEGIN TRANSACTION;
    # CREATE TABLE raw_packets (timestamp NUMERIC, appRawData BLOB, inetRawData BLOB, xportRawData BLOB, recID INTEGER PRIMARY KEY, dstPort NUMERIC, inetProtocol NUMERIC, srcPort NUMERIC, xportProtocol NUMERIC, dstIP TEXT, srcIP TEXT);
    # CREATE INDEX dstIPIdx ON raw_packets(dstIP ASC);
    # CREATE INDEX dstPortIdx ON raw_packets(dstPort ASC);
    # CREATE INDEX srcIPIdx ON raw_packets(srcIP ASC);
    # CREATE INDEX srcPortIdx ON raw_packets(srcPort ASC);
    # COMMIT;
    def createStorage(self):
        
        self.deleteStorage()
        
        try:
            # copy the template file into the actual storage file
            shutil.copyfile(globalStrings.traceStorageTemplate, self.storageFile)
                        
        except Exception, ex:
            strErr = "Could not create trace storage file.\nThere was an error creating physical storage file %s:\n %s" % (self.storageFile, str(ex))
            logger.log(__name__ + ": " + strErr)
            raise Exception, strErr
           
        self.successful = True
    
    def deleteStorage(self):
        try:
            os.remove(self.storageFile)
        except Exception, ex:
            logger.log(__name__ + ": Warning -- Could not delete storage file: " + str(ex))
        
        
        
    ## Opens the storage file    
    def open(self): 
        try:
            self.conn = sqlite.connect(self.storageFile)
            self.cur =  self.conn.cursor()
        except Exception, ex:
            logger.log(__name__ + ": Could not open sqlite database")
            logger.log(ex)
            
            raise ex                
        
        
        
    ## Closes the storage file
    def close(self): 
        try:
            if self.cur is not None:
                self.cur.close()
            
            if self.conn is not None:
                self.conn.close()
        except Exception, ex:
            logger.log(__name__ + ": Could not close sqlite database")
            logger.log(ex)
             
            raise ex

        
        
    ## Saves a single decodedPacketDict Object to the database. 
    # @param decodedPacketDict Dictionary
    def store(self, decodedPacketDict):
        
        if self.successful == False:
            logger.log(__name__ + ": Cannot insert data into the storage database, because there was an error during its creation")
            return
        
        # this needs to be removed once the network decoding has been completed
        if decodedPacketDict == None:
            return
    
        # a restriction of SQLite is that all operations must be done in the same thread.
        
        try:
            netHeader   = sqlite.Binary(decodedPacketDict['networkHeader'])
            xportHeader = sqlite.Binary(decodedPacketDict['transportHeader'])
            
            netHeader     = base64.encodestring(netHeader)
            xportHeader   = base64.encodestring(xportHeader)
            
            if decodedPacketDict['payload'] != None:
                payload       = base64.encodestring(decodedPacketDict['payload'])
            else:
                payload = None
            
            
            vals = (self.recID, time.time(), decodedPacketDict['sourceIP'],
                        decodedPacketDict['destIP'], decodedPacketDict['sourcePort'],
                        decodedPacketDict['destPort'], 
                        decodedPacketDict['netProtocol'], decodedPacketDict['transportProtocol'],
                        netHeader, xportHeader, payload)
            
            self.cur.execute("""insert into raw_packets
                                (recID, 
                                 timestamp, 
                                 srcIP, 
                                 dstIP, 
                                 srcPort, 
                                 dstPort, 
                                 inetProtocol, 
                                 xportProtocol, 
                                 inetRawData, 
                                 xportRawData, 
                                 appRawData)
                      values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", vals)
                                
            self.recID += 1
        except Exception, ex:
               logger.log(__name__ + ": There was an error storing packet to database")
               logger.log(ex)
  
    
    
    ## Performs a Commit.
    def commit(self):
        try:
            self.conn.commit()
        except Exception, ex:
               logger.log(__name__ + ": There was an error committing data")
               logger.log(ex)   
               raise ex     
    
    ## Deleted all data from the database file.
    def clear(self):         
        try:
            self.open()
            self.cur.execute("""delete from raw_packets""")
            
            # Save (commit) the changes
            self.conn.commit()
            self.recID = 0
            
            self.close()
        except Exception, ex:
               logger.log(__name__ + ": Could not clear packet database")
               logger.log(ex)
               raise ex   
    
    
    ## Returns the size of the database file in KB     
    def getStorageSizeBytes(self):
        
        try:
            statInfo = os.stat(globalStrings.traceStoragePath)
            fileSze = statInfo.st_size  
            return fileSze
        except Exception, ex:
            logger.log(__name__ + ": There was an error figuring out storage file size")
            logger.log(ex)
            return 0
        else:
            return fileSze
    
    
    
    ## Returns a RecordEet given the supplied parameters
    # @param dstIP Destination IP Address String
    # @param srcPort Source Port Integer
    # @param  dstPort Destination Port Integer
    def query(self, dstIP, srcPort, dstPort):
        
        try:
            
            self.resultSet = []
            
            vals = (dstIP, dstIP, srcPort, srcPort, dstPort, dstPort)
            
            query = """select timestamp,
                                       srcIP,
                                       dstIP,
                                       srcPort,
                                       dstPort,
                                       inetProtocol,
                                       xportProtocol,
                                       inetRawData, xportRawData, appRawData 
                                from raw_packets
                                where (srcIP = ? or dstIP = ?)
                                  and ( (srcPort = ? or srcPort = 0) or (dstPort = ? or dstPort = 0))
                                  and ( (srcPort = ? or srcPort = 0) or (dstPort = ? or dstPort = 0))
                                order by recID asc"""
                                
            print query, vals
            
            self.cur.execute(query, vals)
            
            map(self.__decodeAndAppend, self.cur)
            
            return self.resultSet
        
        except Exception, ex:
               logger.log(__name__ + ": Could not query database")
               logger.log(ex)
               
               raise ex
    
    
    ## Since all data is stored in Base64 format, this function will perform the appropriate decoding
    # @param row Tuple representing a single row of data.
    def __decodeAndAppend(self, row):
        
        ts, srcip, dstip, srcprt, dstprt, inetp, xportp, netHdr, xptHdr, payload = row
        
        netHdr  = base64.decodestring(netHdr)
        xptHdr  = base64.decodestring(xptHdr)
        
        if payload != None:
            payload = base64.decodestring(payload)
        
        self.resultSet.append((ts, srcip, dstip, srcprt, dstprt, inetp, xportp, netHdr, xptHdr, payload))
        
        
                                 
        

---