eulaurarien/database.py

#!/usr/bin/env python3

"""
Utility functions to interact with the database.
"""

import sqlite3
import typing
import time
import os
import logging
import argparse
import coloredlogs
import ipaddress
import math

coloredlogs.install(
    level='DEBUG',
    fmt='%(asctime)s %(name)s %(levelname)s %(message)s'
)

DbValue = typing.Union[None, int, float, str, bytes]


class Database():
    VERSION = 5
    PATH = "blocking.db"

    def open(self) -> None:
        mode = 'rwc' if self.write else 'ro'
        uri = f'file:{self.PATH}?mode={mode}'
        self.conn = sqlite3.connect(uri, uri=True)
        cursor = self.conn.cursor()
        cursor.execute("PRAGMA foreign_keys = ON")
        self.conn.create_function("unpack_asn", 1,
                                  self.unpack_asn,
                                  deterministic=True)
        self.conn.create_function("unpack_ip4address", 1,
                                  self.unpack_ip4address,
                                  deterministic=True)
        self.conn.create_function("unpack_ip4network", 2,
                                  self.unpack_ip4network,
                                  deterministic=True)
        self.conn.create_function("unpack_domain", 1,
                                  lambda s: s[:-1][::-1],
                                  deterministic=True)
        self.conn.create_function("format_zone", 1,
                                  lambda s: '*' + s[::-1],
                                  deterministic=True)

    def get_meta(self, key: str) -> typing.Optional[int]:
        cursor = self.conn.cursor()
        try:
            cursor.execute("SELECT value FROM meta WHERE key=?", (key,))
        except sqlite3.OperationalError:
            return None
        for ver, in cursor:
            return ver
        return None

    def set_meta(self, key: str, val: int) -> None:
        cursor = self.conn.cursor()
        cursor.execute("INSERT INTO meta VALUES (?, ?) "
                       "ON CONFLICT (key) DO "
                       "UPDATE set value=?",
                       (key, val, val))

    def close(self) -> None:
        self.enter_step('close_commit')
        self.conn.commit()
        self.enter_step('close')
        self.conn.close()
        self.profile()

    def initialize(self) -> None:
        self.close()
        self.enter_step('initialize')
        if not self.write:
            self.log.error("Cannot initialize in read-only mode.")
            raise
        os.unlink(self.PATH)
        self.open()
        self.log.info("Creating database version %d.", self.VERSION)
        cursor = self.conn.cursor()
        with open("database_schema.sql", 'r') as db_schema:
            cursor.executescript(db_schema.read())
        self.set_meta('version', self.VERSION)
        self.conn.commit()

    def __init__(self, write: bool = False) -> None:
        self.log = logging.getLogger('db')
        self.time_last = time.perf_counter()
        self.time_step = 'init'
        self.time_dict: typing.Dict[str, float] = dict()
        self.step_dict: typing.Dict[str, int] = dict()
        self.write = write

        self.open()
        version = self.get_meta('version')
        if version != self.VERSION:
            if version is not None:
                self.log.warning(
                    "Outdated database version: %d found, will be rebuilt.",
                    version)
            self.initialize()

    def enter_step(self, name: str) -> None:
        now = time.perf_counter()
        try:
            self.time_dict[self.time_step] += now - self.time_last
            self.step_dict[self.time_step] += 1
        except KeyError:
            self.time_dict[self.time_step] = now - self.time_last
            self.step_dict[self.time_step] = 1
        self.time_step = name
        self.time_last = time.perf_counter()

    def profile(self) -> None:
        self.enter_step('profile')
        total = sum(self.time_dict.values())
        for key, secs in sorted(self.time_dict.items(), key=lambda t: t[1]):
            times = self.step_dict[key]
            self.log.debug(f"{key:<20}: {times:9d} × {secs/times:5.3e} "
                           f"= {secs:9.2f} s ({secs/total:7.2%}) ")
        self.log.debug(f"{'total':<20}:                         "
                       f"{total:9.2f} s ({1:7.2%})")

    @staticmethod
    def pack_hostname(hostname: str) -> str:
        return hostname[::-1] + '.'

    @staticmethod
    def pack_zone(zone: str) -> str:
        return Database.pack_hostname(zone)

    @staticmethod
    def pack_asn(asn: str) -> int:
        asn = asn.upper()
        if asn.startswith('AS'):
            asn = asn[2:]
        return int(asn)

    @staticmethod
    def unpack_asn(asn: int) -> str:
        return f'AS{asn}'

    @staticmethod
    def pack_ip4address(address: str) -> int:
        total = 0
        for i, octet in enumerate(address.split('.')):
            total += int(octet) << (3-i)*8
        if total > 0xFFFFFFFF:
            raise ValueError
        return total
        # return '{:02x}{:02x}{:02x}{:02x}'.format(
        #     *[int(c) for c in address.split('.')])
        # return base64.b16encode(packed).decode()
        # return '{:08b}{:08b}{:08b}{:08b}'.format(
        #     *[int(c) for c in address.split('.')])
        # carg = ctypes.c_wchar_p(address)
        # ret = ACCEL.ip4_flat(carg, self.accel_ip4_buf)
        # if ret != 0:
        #     raise ValueError
        # return self.accel_ip4_buf.value
        # packed = ipaddress.ip_address(address).packed
        # return packed

    @staticmethod
    def unpack_ip4address(address: int) -> str:
        return '.'.join(str((address >> (i * 8)) & 0xFF)
                        for i in reversed(range(4)))

    @staticmethod
    def pack_ip4network(network: str) -> typing.Tuple[int, int]:
        # def pack_ip4network(network: str) -> str:
        net = ipaddress.ip_network(network)
        mini = Database.pack_ip4address(net.network_address.exploded)
        maxi = Database.pack_ip4address(net.broadcast_address.exploded)
        # mini = net.network_address.packed
        # maxi = net.broadcast_address.packed
        return mini, maxi
        # return Database.pack_ip4address(net.network_address.exploded)[:net.prefixlen]

    @staticmethod
    def unpack_ip4network(mini: int, maxi: int) -> str:
        addr = Database.unpack_ip4address(mini)
        prefixlen = 32-int(math.log2(maxi-mini+1))
        return f'{addr}/{prefixlen}'

    def update_references(self) -> None:
        self.enter_step('update_refs')
        cursor = self.conn.cursor()
        cursor.execute('UPDATE rules AS r SET refs='
                       '(SELECT count(*) FROM rules '
                       'WHERE source=r.id)')

    def prune(self, before: int, base_only: bool = False) -> None:
        self.enter_step('prune')
        cursor = self.conn.cursor()
        cmd = 'DELETE FROM rules WHERE updated<?'
        if base_only:
            cmd += ' AND level=0'
        cursor.execute(cmd, (before,))

    def explain(self, entry: int) -> str:
        # Format current
        string = '???'
        cursor = self.conn.cursor()
        cursor.execute(
            'SELECT unpack_asn(val) FROM asn WHERE entry=:entry '
            'UNION '
            'SELECT unpack_domain(val) FROM hostname WHERE entry=:entry '
            'UNION '
            'SELECT format_zone(val) FROM zone WHERE entry=:entry '
            'UNION '
            'SELECT unpack_ip4address(val) FROM ip4address WHERE entry=:entry '
            'UNION '
            'SELECT unpack_ip4network(mini, maxi) '
            'FROM ip4network WHERE entry=:entry ',
            {"entry": entry}
        )
        for val, in cursor:  # only one
            string = str(val)
        string += f' #{entry}'

        # Add source if any
        cursor.execute('SELECT source FROM rules WHERE id=?', (entry,))
        for source, in cursor:
            if source:
                string += f' ← {self.explain(source)}'
        return string

    def export(self,
               first_party_only: bool = False,
               end_chain_only: bool = False,
               explain: bool = False,
               ) -> typing.Iterable[str]:
        selection = 'entry' if explain else 'unpack_domain(val)'
        command = f'SELECT {selection} FROM rules ' \
            'INNER JOIN hostname ON rules.id = hostname.entry'
        restrictions: typing.List[str] = list()
        if first_party_only:
            restrictions.append('rules.first_party = 1')
        if end_chain_only:
            restrictions.append('rules.refs = 0')
        if restrictions:
            command += ' WHERE ' + ' AND '.join(restrictions)
        if not explain:
            command += ' ORDER BY unpack_domain(val) ASC'
        cursor = self.conn.cursor()
        cursor.execute(command)
        for val, in cursor:
            if explain:
                yield self.explain(val)
            else:
                yield val

    def count_rules(self,
                    first_party_only: bool = False,
                    ) -> str:
        counts: typing.List[str] = list()
        cursor = self.conn.cursor()
        for table in ['asn', 'ip4network', 'ip4address', 'zone', 'hostname']:
            command = f'SELECT count(*) FROM rules ' \
                f'INNER JOIN {table} ON rules.id = {table}.entry ' \
                'WHERE rules.level = 0'
            if first_party_only:
                command += ' AND first_party=1'
            cursor.execute(command)
            count, = cursor.fetchone()
            if count > 0:
                counts.append(f'{table}: {count}')

        return ', '.join(counts)

    def get_domain(self, domain: str) -> typing.Iterable[int]:
        self.enter_step('get_domain_prepare')
        domain_prep = self.pack_hostname(domain)
        cursor = self.conn.cursor()
        self.enter_step('get_domain_select')
        cursor.execute(
            'SELECT null, entry FROM hostname '
            'WHERE val=:d '
            'UNION '
            'SELECT * FROM ('
            'SELECT val, entry FROM zone '
            # 'WHERE val>=:d '
            # 'ORDER BY val ASC LIMIT 1'
            'WHERE val<=:d '
            'AND instr(:d, val) = 1'
            ')',
            {'d': domain_prep}
        )
        for val, entry in cursor:
            # print(293, val, entry)
            self.enter_step('get_domain_confirm')
            if not (val is None or domain_prep.startswith(val)):
                # print(297)
                continue
            self.enter_step('get_domain_yield')
            yield entry

    def get_domain_in_zone(self, domain: str) -> typing.Iterable[int]:
        self.enter_step('get_domainiz_prepare')
        domain_prep = self.pack_hostname(domain)
        cursor = self.conn.cursor()
        self.enter_step('get_domainiz_select')
        cursor.execute(
            'SELECT val, entry FROM zone '
            'WHERE val<=:d '
            'ORDER BY val DESC LIMIT 1',
            {'d': domain_prep}
        )
        for val, entry in cursor:
            self.enter_step('get_domainiz_confirm')
            if not (val is None or domain_prep.startswith(val)):
                continue
            self.enter_step('get_domainiz_yield')
            yield entry

    def get_ip4(self, address: str) -> typing.Iterable[int]:
        self.enter_step('get_ip4_prepare')
        try:
            address_prep = self.pack_ip4address(address)
        except (ValueError, IndexError):
            self.log.error("Invalid ip4address: %s", address)
            return
        cursor = self.conn.cursor()
        self.enter_step('get_ip4_select')
        cursor.execute(
            'SELECT entry FROM ip4address '
            # 'SELECT null, entry FROM ip4address '
            'WHERE val=:a '
            'UNION '
            # 'SELECT * FROM ('
            # 'SELECT val, entry FROM ip4network '
            # 'WHERE val<=:a '
            # 'AND instr(:a, val) > 0 '
            # 'ORDER BY val DESC'
            # ')'
            'SELECT entry FROM ip4network '
            'WHERE :a BETWEEN mini AND maxi ',
            {'a': address_prep}
        )
        for entry, in cursor:
            # self.enter_step('get_ip4_confirm')
            # if not (val is None or val.startswith(address_prep)):
            #     # PERF startswith but from the end
            #     continue
            self.enter_step('get_ip4_yield')
            yield entry

    def get_ip4_in_network(self, address: str) -> typing.Iterable[int]:
        self.enter_step('get_ip4in_prepare')
        try:
            address_prep = self.pack_ip4address(address)
        except (ValueError, IndexError):
            self.log.error("Invalid ip4address: %s", address)
            return
        cursor = self.conn.cursor()
        self.enter_step('get_ip4in_select')
        cursor.execute(
            'SELECT entry FROM ip4network '
            'WHERE :a BETWEEN mini AND maxi ',
            {'a': address_prep}
        )
        for entry, in cursor:
            self.enter_step('get_ip4in_yield')
            yield entry

    def list_asn(self) -> typing.Iterable[typing.Tuple[str, int]]:
        cursor = self.conn.cursor()
        self.enter_step('list_asn_select')
        cursor.execute('SELECT val, entry FROM asn')
        for val, entry in cursor:
            yield f'AS{val}', entry

    def _set_generic(self,
                     table: str,
                     select_query: str,
                     insert_query: str,
                     prep: typing.Dict[str, DbValue],
                     updated: int,
                     is_first_party: bool = False,
                     source: int = None,
                     ) -> None:
        # Since this isn't the bulk of the processing,
        # here abstraction > performaces

        # Fields based on the source
        self.enter_step(f'set_{table}_prepare')
        cursor = self.conn.cursor()
        if source is None:
            first_party = int(is_first_party)
            level = 0
        else:
            self.enter_step(f'set_{table}_source')
            cursor.execute(
                'SELECT first_party, level FROM rules '
                'WHERE id=?',
                (source,)
            )
            first_party, level = cursor.fetchone()
            level += 1

        self.enter_step(f'set_{table}_select')
        cursor.execute(select_query, prep)

        rules_prep: typing.Dict[str, DbValue] = {
            "source": source,
            "updated": updated,
            "first_party": first_party,
            "level": level,
        }

        # If the entry already exists
        for entry, in cursor:  # only one
            self.enter_step(f'set_{table}_update')
            rules_prep['entry'] = entry
            cursor.execute(
                'UPDATE rules SET '
                'source=:source, updated=:updated, '
                'first_party=:first_party, level=:level '
                'WHERE id=:entry AND (updated<:updated OR '
                'first_party<:first_party OR level<:level)',
                rules_prep
            )
            # Only update if any of the following:
            # - the entry is outdataed
            # - the entry was not a first_party but this is
            # - this is closer to the original rule
            return

        # If it does not exist

        self.enter_step(f'set_{table}_insert')
        cursor.execute(
            'INSERT INTO rules '
            '(source, updated, first_party, level) '
            'VALUES (:source, :updated, :first_party, :level) ',
            rules_prep
        )
        cursor.execute('SELECT id FROM rules WHERE rowid=?',
                       (cursor.lastrowid,))
        for entry, in cursor:  # only one
            prep['entry'] = entry
            cursor.execute(insert_query, prep)
            return
        assert False

    def set_hostname(self, hostname: str,
                     *args: typing.Any, **kwargs: typing.Any) -> None:
        self.enter_step('set_hostname_prepare')
        prep: typing.Dict[str, DbValue] = {
            'val': self.pack_hostname(hostname),
        }
        self._set_generic(
            'hostname',
            'SELECT entry FROM hostname WHERE val=:val',
            'INSERT INTO hostname (val, entry) '
            'VALUES (:val, :entry)',
            prep,
            *args, **kwargs
        )

    def set_asn(self, asn: str,
                *args: typing.Any, **kwargs: typing.Any) -> None:
        self.enter_step('set_asn_prepare')
        try:
            asn_prep = self.pack_asn(asn)
        except ValueError:
            self.log.error("Invalid asn: %s", asn)
            return
        prep: typing.Dict[str, DbValue] = {
            'val': asn_prep,
        }
        self._set_generic(
            'asn',
            'SELECT entry FROM asn WHERE val=:val',
            'INSERT INTO asn (val, entry) '
            'VALUES (:val, :entry)',
            prep,
            *args, **kwargs
        )

    def set_ip4address(self, ip4address: str,
                       *args: typing.Any, **kwargs: typing.Any) -> None:
        self.enter_step('set_ip4add_prepare')
        try:
            ip4address_prep = self.pack_ip4address(ip4address)
        except (ValueError, IndexError):
            self.log.error("Invalid ip4address: %s", ip4address)
            return
        prep: typing.Dict[str, DbValue] = {
            'val': ip4address_prep,
        }
        self._set_generic(
            'ip4add',
            'SELECT entry FROM ip4address WHERE val=:val',
            'INSERT INTO ip4address (val, entry) '
            'VALUES (:val, :entry)',
            prep,
            *args, **kwargs
        )

    def set_zone(self, zone: str,
                 *args: typing.Any, **kwargs: typing.Any) -> None:
        self.enter_step('set_zone_prepare')
        prep: typing.Dict[str, DbValue] = {
            'val': self.pack_zone(zone),
        }
        self._set_generic(
            'zone',
            'SELECT entry FROM zone WHERE val=:val',
            'INSERT INTO zone (val, entry) '
            'VALUES (:val, :entry)',
            prep,
            *args, **kwargs
        )

    def set_ip4network(self, ip4network: str,
                       *args: typing.Any, **kwargs: typing.Any) -> None:
        self.enter_step('set_ip4net_prepare')
        try:
            ip4network_prep = self.pack_ip4network(ip4network)
        except (ValueError, IndexError):
            self.log.error("Invalid ip4network: %s", ip4network)
            return
        prep: typing.Dict[str, DbValue] = {
            'mini': ip4network_prep[0],
            'maxi': ip4network_prep[1],
        }
        self._set_generic(
            'ip4net',
            'SELECT entry FROM ip4network WHERE mini=:mini AND maxi=:maxi',
            'INSERT INTO ip4network (mini, maxi, entry) '
            'VALUES (:mini, :maxi, :entry)',
            prep,
            *args, **kwargs
        )


if __name__ == '__main__':

    # Parsing arguments
    parser = argparse.ArgumentParser(
        description="Database operations")
    parser.add_argument(
        '-i', '--initialize', action='store_true',
        help="Reconstruct the whole database")
    parser.add_argument(
        '-p', '--prune', action='store_true',
        help="Remove old entries from database")
    parser.add_argument(
        '-b', '--prune-base', action='store_true',
        help="TODO")
    parser.add_argument(
        '-s', '--prune-before', type=int,
        default=(int(time.time()) - 60*60*24*31*6),
        help="TODO")
    parser.add_argument(
        '-r', '--references', action='store_true',
        help="Update the reference count")
    args = parser.parse_args()

    DB = Database(write=True)

    if args.initialize:
        DB.initialize()
    if args.prune:
        DB.prune(before=args.prune_before, base_only=args.prune_base)
    if args.references:
        DB.update_references()

    DB.close()