MongoUtils.py

#!/usr/bin/env python
#
# MongoUtils.py
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# 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, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston,
# MA  02111-1307  USA

import sys
import getopt

MONGO_MODULE = True

try:
    import gridfs
    import pymongo
    from bson.objectid import ObjectId
except ImportError:
    MONGO_MODULE = False


class MongoUtils(object):
    def __init__(self, host="localhost", port=27017):
        self.__check_mongo_module()
        self.__init_db(host, port)

    def __check_mongo_module(self):
        if not MONGO_MODULE:
            print("[MongoUtils] PyMongo not installed. Please install it and retry.")
            sys.exit(0)

    def __init_db(self, host, port):
        # MongoDB Connection class is marked as deprecated (MongoDB >= 2.4).
        # The following code tries to use the new MongoClient if available and
        # reverts to the old Connection class if not. This code will hopefully
        # disappear in the next future.
        client = getattr(pymongo, "MongoClient", None)
        if client is None:
            client = getattr(pymongo, "Connection", None)

        try:
            connection = client(host, int(port))
        except Exception:
            print("[MongoUtils] MongoDB instance not available")
            sys.exit(0)

        db = connection.thug
        dbfs = connection.thugfs
        self.urls = db.urls
        self.analyses = db.analyses
        self.thugfs = gridfs.GridFS(dbfs)

        self.collections = (
            db.urls,
            db.analyses,
            db.locations,
            db.connections,
            db.graphs,
            db.samples,
            db.behaviors,
            db.certificates,
            db.honeyagent,
            db.exploits,
            db.codes,
            db.json,
        )

    def list_analyses(self):
        print("ID\t\t\t\t| URL\t\t\t\t| Analysis datetime\n")

        for analysis in self.analyses.find():
            urls = self.urls.find(_id=analysis["url_id"])
            if not urls:
                continue

            url = urls[0]["url"]
            print("%s\t| %s\t| %s" % (analysis["_id"], url, analysis["timestamp"]))

    def remove_analysis(self, analysis):
        analysis_id = analysis["_id"]
        for collection in self.collections:
            collection.remove({"_id": analysis_id})

        self.thugfs.delete({"_id": analysis_id})

    def query_analysis_by_id(self, analysis_id):
        analysis = self.analyses.find_one({"_id": ObjectId(analysis_id)})
        if not analysis:
            print("[MongoUtils] Analysis ID not found")
            return None

        return analysis


def usage():
    msg = """
Synopsis:
    MongoUtils.py

    Usage:
        python MongoUtils.py [ options ]

    Options:
        -h, --help              \tDisplay this help information
        -l, --ls                \tList all the analyses
        -r, --rm=               \tRemove the analysis identified by the specified ID
        -M, --mongodb-address=  \tSpecify address and port of the MongoDB instance (format: host:port)
"""
    print(msg)
    sys.exit(0)


def main(args):
    host = "localhost"
    port = 27017

    try:
        options, args = getopt.getopt(
            args, "hlr:M:", ["help", "ls", "rm=", "mongodb-address="]
        )
    except getopt.GetoptError:
        usage()

    if not options and not args:
        usage()

    for option in options:
        if option[0] in ("-h", "--help"):
            usage()
        elif option[0] in ("-M", "--mongodb-address"):
            host, port = option[1].split(":")

    mongoutils = MongoUtils(host, port)

    for option in options:
        if option[0] in ("-l", "--ls"):
            mongoutils.list_analyses()

    for option in options:
        if option[0] in ("-r", "--rm"):
            analysis = mongoutils.query_analysis_by_id(option[1])
            if analysis:
                mongoutils.remove_analysis(analysis)


if __name__ == "__main__":
    main(sys.argv[1:])