fourier-install/tools/fourier/cli.py

from __future__ import print_function
from argparse import ArgumentParser
from hashlib import md5
from datetime import datetime
import sqlite3
import time
import json
import sys
import os

strptime = datetime.strptime
ARG_DATE_FMT = '%Y-%m-%d %H:%M:%S'
simple_date = lambda s: strptime(s, ARG_DATE_FMT)

def find_hash(hsh, db=None):
    cursor = db.cursor()
    cursor.execute("select * from file where hash = ?",
                   (hsh, ),
                   )
    return cursor.fetchone()

def insert_file(hsh, filename, timestamp, db=None):
    cursor = db.cursor()
    cursor.execute("""
        insert into file(hash, filename, timestamp)
        values(?, ?, ?)
        """,
        (hsh, filename, int(time.mktime(timestamp.timetuple())), )
    )

def hash_file(filename):
    ahash = md5()
    with open(filename, 'rb') as fp:
        while True:
            data = fp.read(4096)
            if data:
                ahash.update(data)
            else:
                break
    return ahash.hexdigest()

def main():
    parser = ArgumentParser()
    subparsers = parser.add_subparsers(dest='entity')
    
    dbparser = subparsers.add_parser('db')
    dbparser.add_argument('action')
    dbparser.add_argument('-from-date', dest='from_date',
                                        type=simple_date,
                                        )
    dbparser.add_argument('-station', dest='station')
    dbparser.add_argument('--dry-run', dest='dry_run',
                          action='store_const',
                          const=True,
                          default=False
                          )

    args = parser.parse_args()
    action = args.action

    with open('/etc/fourier-config.json', 'r') as fp:
        config = json.loads(fp.read())

    device_id = config['device_id']
    dbpath = '/var/fourier/{}/files.db'.format(device_id)

    if action == 'stats':
        conn = sqlite3.connect(dbpath)
        cursor = conn.cursor()
        cursor.execute("select count(*), count(uploaded) from file")
        total, uploaded, = cursor.fetchone()
        print("total: {}".format(total))
        print("uploaded: {}".format(uploaded))
        print("pending: {}".format(total - uploaded))

    elif action == 'index-files':
        counter = 0
        already_indexed = 0
        conn = sqlite3.connect(dbpath)
        path = os.path.join('/var/fourier', device_id)
        if args.station:
            path = os.path.join(path, args.station)
        for folder, folders, files in os.walk(path):
            for file in files:
                if not file.endswith('.mp3'):
                    continue

                filename = os.path.join(folder, file)

                dt = datetime.strptime(
                    file[:19],
                    '%Y-%m-%dT%H-%M-%S'
                )
                try:
                    if args.from_date:
                        do_insert = dt >= args.from_date  
                    else:
                        do_insert = True

                    if do_insert:
                        thehash = hash_file(filename)
                        insert_file(thehash, filename, dt, db=conn)
                        counter += 1
                        print(dt)

                except sqlite3.IntegrityError:
                    already_indexed += 1
                    print('already indexed: {}'.format(filename))

        if not args.dry_run:
            conn.commit()
        else:
            conn.rollback()
            print('\n[WARNING] DRY RUN FINISHED')

        print('----------------------------------')
        print('total files indexed: {}'.format(counter))
        print('total files in existence: {}'.format(already_indexed))
        print('----------------------------------')

    elif action == 'setup':
        if not os.path.isfile(dbpath):
            conn = sqlite3.connect(dbpath)
            cursor = conn.cursor()
            sentences = [
                """create table file(
                    hash text primary key,
                    station text,
                    timestamp int,
                    filename text,
                    uploaded int
                )""",
                "create index timestamp_index_desc on file (timestamp desc)",
                "create index timestamp_index_asc on file (timestamp desc)",
            ]
            for query in sentences:
                cursor.execute(query)
            conn.commit()
        else:
            print('database already installed')
            sys.exit(1)
                        

if __name__ == '__main__':
    main()