''' 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