Lark/hashdb.py

'''
hashdb.py
'''
import collections
import hashlib
import sqlite3

# TODO: Decide on a way to auto-download DATs.
# TODO: l o g g i n g
HASH_CHUNK_SIZE = 10485760 # 10mb
SQL_AND = ' AND '
SQL_OR = ' OR '

# TODO: Figure out how to do some kind of type checking for these named tuples.
RomInfo = collections.namedtuple('RomInfo', 'sha1sum, filename, platform, datorigin')
DatInfo = collections.namedtuple('DatInfo', 'name, description, platform, version')
PlatformInfo = collections.namedtuple('PlatformInfo', 'shortcode, fullname, aliases')

ORPHAN_DAT = DatInfo('', 'Orphaned hashes', 'nonexistent', '1')

# TODO: This should go in the eventual romdb class.
def get_file_sha1sum(filename):
    sha1sum = hashlib.sha1()
    with open(filename, 'rb') as file_contents:
        while True:
            chunk = file_contents.read(HASH_CHUNK_SIZE)
            if not chunk:
                break
            sha1sum.update(chunk)

    return sha1sum.hexdigest()

def _build_sql_constraints(inclusive, constraints):
    '''Build an SQL constraint clause out of a dictionary.

    inclusive   - If True, uses SQL AND operator, otherwise OR will be used.
    constraints - A dictionary of arbitrary constraints

    returns a tuple containing the appropriately formatted SQL string and a list of parameters
    '''
    if constraints == {}:
        return ('', [])

    if inclusive:
        logical_separator = SQL_AND
    else:
        logical_separator = SQL_OR

    sql_constraint_string = 'WHERE '
    sql_parameter_list = []
    for key, value in constraints.items():
        sql_constraint_string += '%s=?%s' % (key, logical_separator)
        sql_parameter_list.append(value)

    # Trim off the last ', '
    sql_constraint_string = sql_constraint_string[0:-len(logical_separator)]

    return (sql_constraint_string, sql_parameter_list)

class HashDB:
    '''Store and retrieve hash metadata from an SQLite database.'''
    # TODO: Low-priority: Probably design this around using multiple hash algorithms eventually.
    def __init__(self, filename):
        '''
        If db file does not exist, create it and create necessary tables.
        Either way, create a connection object.
        '''
        # TODO: This process needs real error handling.
        self._connection = sqlite3.connect(filename)

        with self._connection:
            # TODO: sha1sums.datorigin should be treated as a list. 
            self._connection.execute('CREATE TABLE IF NOT EXISTS sha1sums (sha1sum PRIMARY KEY, '
                                     'filename NOT NULL, platform NOT NULL, datorigin);')

            # TODO: Consider moving image-dat association to dats table.
            self._connection.execute('CREATE TABLE IF NOT EXISTS dats (name PRIMARY KEY, '
                                     'description, platform NOT NULL, version NOT NULL);')

            # TODO: Add support for custom roms not tracked in DAT releases.
            # INSERT INTO dats (name="custom", description="Personally added hashes.", version=1);

            self._connection.execute('CREATE TABLE IF NOT EXISTS platforms (shortcode PRIMARY KEY, '
                                     'fullname NOT NULL, aliases );')
            print('Database initialized.')


    def add_hash(self, rom_info):
        ''' Add a hash to the database. '''
        # INSERT INTO sha1sums (sha1sum, filename, platform, datorigin);
        with self._connection:
            self._connection.execute('INSERT INTO sha1sums VALUES (?, ?, ?, ?)', rom_info)

    def add_hash_list(self, rom_info_list):
        '''Add many hashes to the database. '''
        with self._connection:
            for rom_info in rom_info_list:
                self._connection.execute('INSERT INTO sha1sums VALUES (?, ?, ?, ?)', rom_info)

    def remove_hash(self, rom_info):
        ''' Remove a hash from the database. '''
        # DELETE FROM sha1sums WHERE sha1sum=sha1sum;
        with self._connection:
            self._connection.execute('DELETE FROM sha1sums WHERE sha1sum=?;', [rom_info.sha1sum])

    def remove_hash_list(self, rom_info_list):
        '''Remove many hashes from the database. '''
        with self._connection:
            for rom_info in rom_info_list:
                self._connection.execute('DELETE FROM sha1sums WHERE sha1sum=?;', [rom_info.sha1sum])

    def add_platform(self, platform_info):
        ''' Add a platform shortcode to the database. '''
        # TODO: Collisions need user input to resolve, so remove this try block later.
        try:
            with self._connection:
                self._connection.execute('INSERT INTO platforms VALUES (?, ?, ?);', platform_info)
        except sqlite3.IntegrityError:
            print('Warning: %s is already in database.' % platform_info.shortcode)

    def update_platform_aliases(self, shortcode, aliases):
        ''' Change the list of aliases for a platform shortcode '''
        # UPDATE platforms SET aliases=aliases WHERE shortcode=shortcode;

    def remove_platform(self, platform_info):
        ''' Remove a platform and all associated DATs and hashes from the database. '''
        # DELETE FROM sha1sums WHERE platform=shortcode;
        # DELETE FROM dats WHERE platform=shortcode;
        # DELETE FROM platform WHERE platform=shortcode;
        with self._connection:
            self._connection.execute('DELETE FROM sha1sums WHERE platform=?;',
                                     [platform_info.shortcode])
            self._connection.execute('DELETE FROM dats WHERE platform=?;',
                                     [platform_info.shortcode])
            self._connection.execute('DELETE FROM platforms WHERE shortcode=?;',
                                     [platform_info.shortcode])

    def add_dat(self, dat_info):
        '''Add a DAT's metadata to the database. '''
        with self._connection:
            self._connection.execute('INSERT INTO platforms VALUES (?, ?, ?, ?);', dat_info)

    def remove_dat(self, dat_info):
        ''' Delete a DAT and all of its' hashes from the database. '''
        # DELETE FROM sha1sums WHERE datorigin=name;
        # DELETE FROM dats WHERE name=name;

        with self._connection:
            # TODO: Support multiple dat sources for the same hash.
            self._connection.execute('DELETE FROM sha1sums WHERE datorigin=?;', [dat_info.name])
            self._connection.execute('DELETE FROM dats WHERE name=?;', [dat_info.name])

    def hash_search(self, inclusive=True, **constraints):
        '''Search for hashes, given the parameters. '''

        sql_where_clause, sql_parameters = _build_sql_constraints(inclusive, constraints)

        rom_info_list = []
        with self._connection:
            cursor = self._connection.cursor()
            sql_query = 'SELECT * FROM sha1sums %s;' % sql_where_clause
            cursor.execute(sql_query, sql_parameters)
            print(sql_query)
            rows = cursor.fetchall()

            for row in rows:
                rom_info = RomInfo(*row)
                rom_info_list.append(rom_info)

        return rom_info_list