flashair-util

#!/usr/bin/env python3
#coding: UTF-8

import logging
import re
import time
import tabulate

from argparse import ArgumentParser
from functools import partial
from pathlib import Path

from requests import RequestException
from tfatool import command, sync, info, cgi, util


logger = logging.getLogger("main")
logging.basicConfig(level=logging.INFO, style="{",
                    format="{asctime} | {levelname} | {name} | {message}")
logging.getLogger("requests").setLevel(logging.WARNING)

parser = ArgumentParser()
parser.add_argument("-v", "--verbose", action="store_true")

actions = parser.add_argument_group("Actions")
actions.add_argument("-l", "--list-files", action="store_true")
actions.add_argument("-c", "--count-files", action="store_true")
actions.add_argument("-s", "--sync-forever", action="store_true",
                     help="watch for new files in REMOTE_DIR, copy them to LOCAL_DIR "
                          "(runs until CTRL-C)")
actions.add_argument("-S", "--sync-once", default=False,
                     choices=["time", "name", "all"],
                     help="move files (all or by most recent name/timestamp) from "
                          "REMOTE_DIR to LOCAL_DIR, then quit")

setup = parser.add_argument_group("Setup")
setup.add_argument("-y", "--sync-direction", choices=["up", "down", "both"],
                   help="'up' to upload, 'down' to download, 'both' for both "
                        "(default: down)",
                   default="down")
setup.add_argument("-r", "--remote-dir", default=info.DEFAULT_REMOTE_DIR,
                   help="FlashAir directory to work with (default: {})".format(
                        info.DEFAULT_REMOTE_DIR))
setup.add_argument("-d", "--local-dir", default=".",
                   help="local directory to work with (default: working dir)")

filt = parser.add_argument_group("File filters")
filt.add_argument("-j", "--only-jpg", action="store_true",
                   help="filter for only JPEG files")
filt.add_argument("-n","--n-files", type=int, default=1,
                   help="Number of files to move in --sync-once mode")
filt.add_argument("-k", "--match-regex", default=None,
                   help="filter for files that match the given pattern")
filt.add_argument("-t", "--earliest-date",
                   help="work on only files AFTER datetime of any "
                        "reasonable format such as YYYY, YYYY-MM, "
                        "MM/DD, HH:mm (today), or 'YYYY-MM-DD HH:mm:ss'")
filt.add_argument("-T", "--latest-date",
                   help="work on only files BEFORE given datetime")


def run():
    args = parser.parse_args()
    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    # filename filters
    filters = []
    if args.only_jpg:
        jpg_filter = lambda f: f.filename.lower().endswith(".jpg")
        filters.append(jpg_filter)
    if args.match_regex:
        regex_filter = lambda f: re.match(args.match_regex, f.filename)
        filters.append(regex_filter)

    # datetme filters
    if args.earliest_date:
        try:
            earliest_date = util.parse_datetime(args.earliest_date)
            filters.append(lambda f: f.datetime > earliest_date)
        except ValueError as e:
            parser.error("Invalid earliest date: {}".format(str(e)))
    if args.latest_date:
        try:
            latest_date = util.parse_datetime(args.latest_date)
            filters.append(lambda f: f.datetime < latest_date)
        except ValueError as e:
            parser.error("Invalid latest date: {}".format(str(e)))

    if args.earliest_date or args.latest_date:
        early = "beginning of time"
        late = "end of time"
        if args.earliest_date:
            early = earliest_date.format("YYYY-MM-DD HH:mm:ss")
        if args.latest_date:
            late = latest_date.format("YYYY-MM-DD HH:mm:ss")
        logger.info("Filtering from [{}] to [{}]".format(early, late))

    try:
        if args.list_files:
            print_file_list(filters, args)
        if args.count_files:
            print_file_count(filters, args)
    except RequestException as e:
        print("\nHTTP request exception: {}".format(e))

    if args.sync_once == "all" and args.n_files != 1:
        parser.error("`--sync-once all` doesn't make sense with `--num-files N`")

    if args.sync_forever or args.sync_once:
        local_path = Path(args.local_dir)
        if not local_path.is_dir():
            logger.info("Creating directory '{}'".format(args.local_dir))
            local_path.mkdir()
        if args.sync_once:
            methods = iter_sync_once_methods(args)
            sync_once(methods, filters, args)
        if args.sync_forever:
            try:
                sync_loop(filters, args)
            except KeyboardInterrupt:
                pass


def iter_sync_once_methods(args):
    if args.sync_direction in ("up", "both"):
        if args.sync_once == "name":
            yield sync.up_by_name
        elif args.sync_once == "time":
            yield sync.up_by_time
        elif args.sync_once == "all":
            yield sync.up_by_all
    if args.sync_direction in ("down", "both"):
        if args.sync_once == "name":
            yield sync.down_by_name
        elif args.sync_once == "time":
            yield sync.down_by_time
        elif args.sync_once == "all":
            yield sync.down_by_all


def sync_once(methods, filters, args):
    for by_method in methods:
        try:
            by_method(*filters, remote_dir=args.remote_dir,
                      local_dir=args.local_dir, count=args.n_files)
        except KeyboardInterrupt:
            break


def sync_loop(filters, args):
    try:
        _sync_loop(filters, args)
    except KeyboardInterrupt:
        pass


def _sync_loop(filters, args):
    if args.sync_direction == "both":
        run_method = sync.up_down_by_arrival
    elif args.sync_direction == "up":
        run_method = sync.up_by_arrival
    else:
        run_method = sync.down_by_arrival

    while True:
        new_files = run_method(
            *filters, local_dir=args.local_dir,
            remote_dir=args.remote_dir)
        logger.info("Waiting for newly arrived files...")
        try:
            while True:
                _, new = next(new_files)
                if not new:
                    time.sleep(0.3)
        except KeyboardInterrupt:
            break
        else:
            logger.warning("Sync loop interrupted -- restarting in 5s")
            time.sleep(5)


_fields = ["filename", "date", "time", "MB", "created"]


def print_file_list(filters, args, count_only=False):
    files = command.list_files(*filters, remote_dir=args.remote_dir)
    files = list(files)
    rows = util.fmt_file_rows(files)
    table = tabulate.tabulate(rows, headers=_fields, tablefmt="simple")
    print("\nFiles in {}".format(args.remote_dir))
    if not count_only:
        print()
        print(table)
        print()
    nbytes = sum(f.size for f in files)
    size, units = util.get_size_units(nbytes)
    print("({:d} files, {:0.2f} {} total)".format(
          len(files), size, units))
    

print_file_count = partial(print_file_list, count_only=True)


if __name__ == "__main__":
    with cgi.session:
        run()