blockinator/sqlite.py

#!/usr/bin/python
"""Script to create a SQLite DB from the blocklist data."""

import blocklist
import sys
import logging

try:
    import sqlite3
except ImportError:
    from pysqlite2 import dbapi2 as sqlite3

class SQLiteBlockList(blocklist.BlockList):
    """New class to extend the main BlockList class for
       implementation with Varnish."""
    def export(self):
        """Exports blocklist criteria to a SQLite file."""

        try:
            database = sqlite3.connect(self.config.get("sqlite", "database"))
            cur = database.cursor()

            cur.execute(
                "CREATE TABLE IF NOT EXISTS blocklist (" +
                "remote_ip VARCHAR(15)," +
                "forwarded_ip VARCHAR(15)," +
                "useragent VARCHAR(256)," +
                "cookie VARCHAR(1024)," +
                "PRIMARY KEY(remote_ip))"
            )

            cur.execute("DELETE FROM blocklist")

            for item in self.data.values():
                if item["useragent"] == 'NULL':
                    item["useragent"] = '(null)'

                cur.execute(
                    "INSERT INTO blocklist VALUES ('%s', '%s', '%s', '%s')" % (
                        item["remote_ip"],
                        item["forwarded_ip"],
                        item["useragent"],
                        item["cookie"]
                    )
                )
            database.commit()
            cur.close()
        except sqlite3.Error, error:
            logging.error(
                "There was a problem exporting the data to SQLite: %s",
                error
            )

def main():
    """Main program loop."""
    block_list = SQLiteBlockList()

    block_list.read("http://example.com/default/blocks.txt")

    block_list.export()

    sys.exit(0)

if __name__ == "__main__":
    main()