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
                          )

    stations_parser = subparsers.add_parser('station')
    stations_parser.add_argument('action')

    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 args.entity == 'db':
        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)

    elif args.entity == 'station':
        if action == 'list':
            stations_path = os.path.join('/var/fourier', device_id)
            dirs = os.listdir(stations_path)
            for dr in dirs:
                if '.' not in dr:
                    print(dr)
                        

if __name__ == '__main__':
    main()