oatmeal.py

#!/usr/bin/env python
"""
Oatmeal. A simple cookie manager for Chromium-based browser engines.

Copyright (C) 2021, Judd Vinet <jvinet@zeroflux.org>.

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""
import sys
import os
from os.path import isfile
import sqlite3
import json
import re
import readline
import argparse
from dataclasses import dataclass
from datetime import datetime, timezone
from functools import partial, lru_cache

from rich import box
from rich.console import Console
from rich.table import Table
from rich.theme import Theme

__version__ = '0.1'

DB_PATH = './Cookies'
WL_PATH = './whitelist.json'
BL_PATH = './blacklist.json'

interactive = os.isatty(sys.stdin.fileno())

whitelist = []
blacklist = []

@dataclass()
class Selection:
    """
    This is the list of "selected" cookies that will be used or mutated by
    commands. The selection list is modified by the view commands (eg, 'ls').
    """
    data: list
    title: str
    page: int
    type: str

    def set(self, data, title, page=0, type='cookies'):
        self.data = data
        self.title = title
        self.page = page
        self.type = type  # cookies/whitelist/blacklist

selection = Selection(data=None, title=None, page=0, type=None)

theme = Theme({'info': 'deep_sky_blue1',
               'action': 'green',
               'key': 'pale_turquoise4',
               'value': 'cyan3',
               'error': 'bold red',
               'warning': 'yellow',
               'hl': 'dark_slate_gray1',
               'repr.str': 'bold deep_sky_blue1',
               'repr.tag_name': 'bold yellow',
               'repr.number': 'bold yellow',
               'table.row1': 'cyan',
               'table.row2': 'pale_turquoise4'})
console = Console(theme=theme)


#-----------------------------------------------------------------------------
# COMMANDS
#-----------------------------------------------------------------------------

def cmd_help(command=None, *args):
    """
    Provide usage documentation for commands.

    Usage: h [[repr.tag_name]command[/repr.tag_name]]
    """
    if command:
        if command not in COMMANDS:
            err("Unknown command. Try 'h' for help.")
            return 1
        func = COMMANDS[command]
        if not getattr(func, '__doc__'):
            err(f'No help available for [hl]{command}[/hl]')
            return 0
        info(re.sub(r'Usage: (\w+)', 'Usage: [hl]\\1[/hl]', func.__doc__))
        return 0

    # List all commands.
    table = Table.grid(padding=(0, 4, 0, 0))
    table.add_column(min_width=10, style='info')
    table.add_column(style='white')
    for cmd, func in COMMANDS.items():
        # We have sub-headers packed into the command list, identified by
        # leading octothorpes.
        if cmd.startswith('#'):
            table.add_row(f'  [yellow]{func}[/yellow]', style='bold')
            continue
        # We also have command aliases, which do not appear in the help.
        if cmd.startswith('>'):
            continue
        if getattr(func, '__doc__'):
            doc = func.__doc__.strip().split('\n')[0]
        else:
            doc = ''
        table.add_row(f'    {cmd}', doc)

    console.print('[bold]Commands:[/bold]')
    console.print(table)
    return 0


def cmd_quit(*args):
    """
    Save black/white lists and exit.

    Usage: q
    """
    shutdown()


def cmd_exit(*args):
    """
    Exit without saving black/white lists.

    Usage: x
    """
    shutdown(save=False)


def cmd_list(*args):
    """
    List the current page of the selection.

    Usage: l
    """
    if selection.data is None:
        return warn("No selection.")
    if selection.type == 'cookies':
        return cookie_list()
    return bw_list()


def cmd_count(host=None, *args):
    """
    Count all entries in the selection.

    Usage: c
    """
    ct = len(selection.data)
    s = 'cookies' if selection.type == 'cookies' else f'{selection.type} entries'
    info(f"There are [yellow]{ct}[/yellow] {s}.")


def cmd_next_page(*args):
    """
    Next Page.

    Usage: n
    """
    if (selection.page + 1) * page_size() >= len(selection.data):
        warn("Last page reached.")
        return 0
    selection.page += 1
    if interactive:
        return cmd_list()


def cmd_prev_page(*args):
    """
    Previous page.

    Usage: p
    """
    if selection.page == 0:
        warn("First page reached.")
        return 0
    selection.page -= 1
    if interactive:
        return cmd_list()


def cmd_view(num=None, *args):
    """
    View details for a cookie, referenced by number.

    Usage: v <num>
    """
    if not num:
        return err_arg('num')
    try:
        num = int(num)
    except:
        return err('Invalid number.')
    if num < 1 or num > len(selection.data):
        return err('Number exceeds the selection range.')

    if selection.type != 'cookies':
        warn("This command only works on cookie selections. "
             "Try using 'sca' or 'sch' first.")
        return 1

    # Our list is zero-indexed, but the numbers passed into the command are
    # one-indexed, so we account for that.
    c = selection.data[num - 1]

    table = Table.grid()
    table.add_column(min_width=14, style='key')
    table.add_column(style='value')
    table.add_row('Host:', c['host_key'])
    table.add_row('Name:', c['name'])
    table.add_row('Path:', c['path'])
    table.add_row('Created:', ts_to_datetime(c['creation_utc']))
    table.add_row('Last Used:', ts_to_datetime(c['last_access_utc']))
    table.add_row('Expires:', ts_to_datetime(c['expires_utc']))
    table.add_row('Secure:', str(c['is_secure']))
    table.add_row('HTTP Only:', str(c['is_httponly']))
    table.add_row('Persistent:', str(c['is_persistent']))
    table.add_row('Priority:', str(c['priority']))
    table.add_row('Same Site:', str(c['samesite']))
    table.add_row('Value:', str(c['value']))

    console.print(table)


def cmd_select_cookies_all(*args):
    """
    Select all cookies in DB.

    Usage: sca
    """
    cookies = cookie_find()
    selection.set(cookies, 'All Cookies')


def cmd_select_cookies_by_host(host=None, *args):
    """
    Select all cookies matching host.

    Usage: sch <host>
    """
    if not host:
        return err_arg('host')

    cookies = cookie_find('host_key LIKE ?', (f'%{host}%',))
    selection.set(cookies, f"Cookies for '{host}'")


def cmd_select_blacklist_all(*args):
    """
    Select all entries in blacklist.

    Usage: sba
    """
    selection.set(blacklist, 'Full Blacklist', type='blacklist')


def cmd_select_blacklist_by_host(host=None, *args):
    """
    Select all blacklist entries matching host.

    Usage: sbh <host>
    """
    if not host:
        return err_arg('host')

    # Look for substring matches.
    matches = [x for x in blacklist if host in x]
    selection.set(matches, f"Blacklist Entries for '{host}'", type='blacklist')


def cmd_select_whitelist_all(*args):
    """
    Select all entries in whitelist.

    Usage: swa
    """
    selection.set(whitelist, 'Full Whitelist', type='whitelist')


def cmd_select_whitelist_by_host(host=None, *args):
    """
    Select all whitelist entries matching host.

    Usage: swh <host>
    """
    if not host:
        return err_arg('host')

    # Look for substring matches.
    matches = [x for x in whitelist if host in x]
    selection.set(matches, f"Whitelist Entries for '{host}'", type='whitelist')


def cmd_add_host(host=None, *args):
    """
    Add a host to the selection.

    Usage: ah
    """
    if not host:
        return err_arg('host')
    if selection.type == 'cookies':
        warn("This command only works on the blacklist or whitelist. "
             "Select the blacklist ('sba') or whitelist ('swa') first.")
        return 1

    lst = {'blacklist': blacklist, 'whitelist': whitelist}[selection.type]
    if host in lst:
        err(f"Host '{host}' is already in {selection.type}.")
        return 1

    lst.append(host)
    info(f"Added '{host}' to {selection.type}.")


def cmd_delete_by_number(number_range=None, *args):
    """
    Delete one or more entries by number.

    Usage: dn <number_range>

    Examples:
      dn 3
      dn 6-12
      dn 6,9,3,11,14-17
    """
    if not number_range:
        return err_arg('number_range')

    idxs = []
    # Validate all specifiers before deleting a single record.
    specs = number_range.split(',')
    for spec in specs:
        if not re.match('^[0-9-]+$', spec):
            return err("Argument 'number_range' contains invalid characters.")

        if re.match('^[0-9]+-[0-9]+$', spec):
            start, end = map(int, spec.split('-'))
            if start < 1 or end > len(selection.data):
                return err(f"Index '{spec}' is out of range")
            idxs.extend(list(range(start, end+1)))
        elif re.match('^[0-9]+$', spec):
            spec = int(spec)
            if spec < 1 or spec > len(selection.data):
                return err(f"Index '{spec}' is out of range")
            idxs.extend([spec])
        else:
            return err(f"Number range invalid or unsupported: '{spec}'")

    # Remove them from the DB.
    if selection.type == 'cookies':
        to_del = [selection.data[i - 1] for i in idxs]
        cookie_delete(to_del)

    # Remove them from the selection.
    ct = 0
    for idx in sorted(idxs, reverse=True):
        # Our list is zero-indexed, but the ranges passed into the command are
        # one-indexed, so we account for that.
        idx -= 1

        c = selection.data[idx]
        if selection.type == 'cookies':
            act(f"Deleting: {cookie_repr(c)}")
            del selection.data[idx]
        else:
            act(f"Deleting: {bw_repr(c)}")
            del selection.data[idx]

        ct += 1

    s = f'entries from ' if selection.type != 'cookies' else ''
    info(f"Deleted [yellow]{ct}[/yellow] {s}{selection.type}.")


def cmd_delete_all(*args):
    """
    Delete all entries in selection.

    If the selection target is the cookie list, then cookies whose host matches
    a whitelist entry will not be deleted.

    Usage: da
    """
    idxs = None

    # Remove them from the DB.
    if selection.type == 'cookies':
        idxs =  []
        for idx, c in enumerate(selection.data):
            # Check the whitelist to see if this guy gets a pass.
            if any(host_match(c['host_key'], w) for w in whitelist):
                info(f"Cookie matches whitelist, ignoring: {cookie_repr(c)}")
            else:
                idxs.append(idx)
        cookie_delete([c for i, c in enumerate(selection.data) if i in idxs])

    # Remove them from the selection.
    if idxs is None:
        ct = len(selection.data)
        selection.data[:] = []
    else:
        ct = len(idxs)
        for idx in sorted(idxs, reverse=True):
            del selection.data[idx]

    s = f'entries from ' if selection.type != 'cookies' else ''
    info(f"Deleted [yellow]{ct}[/yellow] {s}{selection.type}.")


def cmd_delete_by_expired(*args):
    """
    Delete expired cookies from the selection.

    Usage: de
    """
    if selection.type != 'cookies':
        warn("This command only works on cookie selections. "
             "Try using 'sca' or 'sch' first.")
        return 1

    ct = 0
    ts = datetime_to_ts(datetime.utcnow())
    idxs = [i for i, c in enumerate(selection.data) if c['expires_utc'] < ts]

    # Remove them from the DB.
    if selection.type == 'cookies':
        to_del = [selection.data[i] for i in idxs]
        cookie_delete(to_del)

    # Remove them from the selection.
    for idx in sorted(idxs, reverse=True):
        del selection.data[idx]
        ct += 1

    info(f"Deleted [yellow]{ct}[/yellow] cookies")


def cmd_delete_by_blacklist(*args):
    """
    Delete all selected cookies whose hosts are found in the blacklist.
    Cookies whose host matches a whitelist entry will not be deleted.

    Usage: db
    """
    if selection.type != 'cookies':
        warn("This command only works on cookie selections. "
             "Try using 'sca' or 'sch' first.")
        return 1

    idxs = []
    for bl_host in blacklist:
        for idx, c in enumerate(selection.data):
            # We'll be passing these indexes to the `cmd_delete_by_number()`
            # command, so these numbers need to be one-indexed.
            idx += 1

            if host_match(c['host_key'], bl_host):
                # Check the whitelist to see if this guy gets a pass.
                if any(host_match(c['host_key'], w) for w in whitelist):
                    info(f"Cookie matches whitelist, ignoring: {cookie_repr(c)}")
                else:
                    idxs.append(idx)

    if idxs:
        return cmd_delete_by_number(','.join(map(str, idxs)))


COMMANDS = {'#1': 'Select',
            'sca': cmd_select_cookies_all,
            'sch': cmd_select_cookies_by_host,
            'sba': cmd_select_blacklist_all,
            'sbh': cmd_select_blacklist_by_host,
            'swa': cmd_select_whitelist_all,
            'swh': cmd_select_whitelist_by_host,

            '#2': 'List',
            'l': cmd_list,
            '>ls': 'l',          # Muscle memory. ¯\_(ツ)_/¯
            'c': cmd_count,
            'n': cmd_next_page,
            'p': cmd_prev_page,

            '#3': 'View',
            'v': cmd_view,

            '#4': 'Add',
            'ah': cmd_add_host,

            '#5': 'Delete',
            'dn': cmd_delete_by_number,
            'da': cmd_delete_all,
            'de': cmd_delete_by_expired,
            'db': cmd_delete_by_blacklist,

            '#6': 'Other',
            'h': cmd_help,
            'x': cmd_exit,
            'q': cmd_quit}


#-----------------------------------------------------------------------------
# BW-LIST MGMT ROUTINES
#-----------------------------------------------------------------------------


def bw_list():
    table = Table(title=selection.title)
    table.add_column('#', style='green', justify='right', no_wrap=True)
    table.add_column('Host', style='cyan3', no_wrap=True)

    idx = 1
    for host in selection.data[:page_size()]:
        table.add_row(str(idx), host)
        idx += 1

    console.print(table)
    return 0


def bw_repr(host):
    """ Return a pretty string-representation of a bw-list entry. """
    return f"[key]host=[/key][value]{host}[/value]"


#-----------------------------------------------------------------------------
# COOKIE MGMT ROUTINES
#-----------------------------------------------------------------------------


def cookie_delete(cookies):
    if not isinstance(cookies, (list, tuple)):
        cookies = [cookies]

    con = sqlite3.connect(DB_PATH)
    for c in cookies:
        con.execute("DELETE FROM cookies WHERE host_key=? AND name=? AND path=?",
                    (c['host_key'], c['name'], c['path']))
    con.commit()
    rc = con.total_changes
    con.close()


def cookie_find(where='1=1', params=()):
    con = sqlite3.connect(DB_PATH)
    con.row_factory = sqlite3.Row

    rv = []
    q = f'SELECT * FROM cookies WHERE {where} ORDER BY creation_utc DESC'
    for row in con.execute(q, params):
        rv.append({k: row[k] for k in row.keys()})

    con.close()
    return rv


def cookie_list():
    table = Table(title=selection.title)
    table.add_column('#', style='green', justify='right', no_wrap=True)
    table.add_column('Host', style='cyan3', no_wrap=True)
    table.add_column('Name')
    table.add_column('Path')
    table.add_column('Created', justify='right', no_wrap=True)
    table.add_column('Expires', justify='right', no_wrap=True)

    pg_sz = page_size()
    start = selection.page * pg_sz
    end = start + pg_sz

    idx = start + 1
    for c in selection.data[start:end]:
        if c['expires_utc']:
            expires = ts_to_datetime(c['expires_utc'])
        else:
            expires = 'none'

        # Shouldn't ever be null, but you never know...
        if c['creation_utc']:
            created = ts_to_datetime(c['creation_utc'])
        else:
            created = 'none'

        s = 'table.row1' if idx % 2 else 'table.row2'
        table.add_row(str(idx), c['host_key'], c['name'], c['path'], created,
                      expires, style=s)
        idx += 1

    console.print(table)


def cookie_repr(c):
    """ Return a pretty string-representation of a cookie. """
    return f"[key]host=[/key][cyan3]{c['host_key']}[/cyan3] " +\
           f"[key]name=[/key][cyan3]{c['name']}[/cyan3] " +\
           f"[key]path=[/key][cyan3]{c['path']}[/cyan3]"


#-----------------------------------------------------------------------------
# MAIN
#-----------------------------------------------------------------------------


def shutdown(save=True):
    if save:
        save_json_file(WL_PATH, whitelist)
        save_json_file(BL_PATH, blacklist)
    # Reclaim deleted space in the DB file.
    con = sqlite3.connect(DB_PATH)
    con.execute("VACUUM")
    con.close()
    sys.exit(0)


def handle_statement(inp):
    inp = inp.strip().split(' ')
    cmd, args = inp[0].strip(), inp[1:]

    if not cmd:
        return 0

    if cmd.startswith(('#', '>')) or cmd not in COMMANDS:
        # Check if it's an alias before rejecting it.
        if f'>{cmd}' in COMMANDS:
            cmd = COMMANDS[f'>{cmd}']
        else:
            err("Unknown command. Try 'h' for help.")
            return 0

    return COMMANDS[cmd](*args) or 0


def process_input(inp):
    # Multiple statements can be specified if they are delimited by semicolons,
    # similar to a real shell.
    for stmt in inp.split(';'):
        rv = handle_statement(stmt)
        # A negative return value indicates that we should abort.
        if rv < 0:
            return rv
    return 0


def loop():
    try:
        inp = console.input('> ')
    except EOFError:
        console.print('q')
        # Tell `main()` to exit.
        return -1
    return process_input(inp)


def main(argv):
    p = argparse.ArgumentParser(description='Manage cookies for QT WebEngine')
    p.add_argument('-c', metavar='PATH', help='location of cookie database file')
    p.add_argument('-b', metavar='PATH', help='location of blacklist file')
    p.add_argument('-w', metavar='PATH', help='location of whitelist file')
    p.add_argument('-e', metavar='CMD',
                   help='execute these commands and exit ' +
                        '(separate multiple commands with semicolons)')
    p.add_argument('-v', action='version', version=f'Oatmeal {__version__}')
    args = p.parse_args(argv[1:])

    if args.c:
        global DB_PATH
        DB_PATH = args.c
    if args.b:
        global BL_PATH
        BL_PATH = args.b
    if args.w:
        global WL_PATH
        WL_PATH = args.w

    # Load whitelist and blacklist.
    global whitelist, blacklist
    try:
        whitelist = load_json_file(WL_PATH)
        blacklist = load_json_file(BL_PATH)
    except:
        sys.exit(1)

    # Ensure the database exists and is actually a database.
    if not isfile(DB_PATH):
        err(f"Database file '{DB_PATH}' not found")
        sys.exit(1)
    try:
        con = sqlite3.connect(DB_PATH)
        con.execute("SELECT * FROM cookies LIMIT 1")
        con.close()
    except sqlite3.DatabaseError as e:
        err("Database error: %s" % e)
        sys.exit(1)

    # If `-e` was used, then we are not in interactive mode. Run the specified
    # commands and bail.
    if args.e:
        global interactive
        interactive = False
        process_input(args.e)
        shutdown()

    # Set & display the initial selection.
    cmd_select_cookies_all()
    cmd_list()

    while True:
        if loop() < 0:
            shutdown()

#-----------------------------------------------------------------------------
# HELPERS
#-----------------------------------------------------------------------------


def load_json_file(path):
    try:
        with open(path, 'r') as f:
            return json.loads(f.read())
    except ValueError:
        err(f"Failed to load '{path}', invalid JSON")
        raise
    except FileNotFoundError:
        return []
    except Exception as e:
        err(f"Failed to read '{path}': {e}")
        raise


def save_json_file(path, body):
    try:
        with open(path, 'w') as f:
            f.write(json.dumps(body))
    except ValueError:
        err(f"Failed to serialize data to JSON")
        raise
    except Exception as e:
        err(f"Failed to write to '{path}': {e}")
        raise
    return None


@lru_cache(maxsize=4096)
def ts_to_datetime(ts):
    # For some presumably-genius reason, the timestamp epoch starts at
    # 1601-01-01T00:00:00Z. And it's in microseconds. You know, for those
    # really high-precision cookies that are necessary to track your site
    # visits down to the millionth of a second since the WWW was first
    # invented in the early 1600s.
    #
    # https://stackoverflow.com/questions/43518199/cookies-expiration-time-format
    ts = ts // 1e6 - 11644473600
    return datetime.fromtimestamp(ts).strftime('%Y-%m-%d %H:%M:%SZ')


@lru_cache(maxsize=4096)
def datetime_to_ts(dt):
    ts = dt.replace(tzinfo=timezone.utc).timestamp()
    # Now transmogrify it into this weirdo version.
    return (ts + 11644473600) * 1e6


@lru_cache(maxsize=4096)
def host_match(c_host, bw_host):
    """ Check if a cookie `c_host` matches a bw-list `bw_host`. """
    if c_host == bw_host:
        return True
    elif bw_host.startswith('.') and c_host.endswith(bw_host):
        return True
    return False


def page_size():
    # Keep only what we can fit on the screen. Leave six rows: three for the
    # table header, one for the title, one for the footer, and one for our
    # command prompt.
    return Console().size.height - 6


def info(msg):
    """ Generate an info message. """
    console.print(msg, style='info')


def act(msg):
    """ Generate a status/action message. """
    console.print(msg, style='action')


def warn(msg):
    """ Generate a warning message. """
    console.print(msg, style='warning')
    return 1


def err(msg):
    """ Generate an error message. """
    console.print(msg, style='error')
    return 1


def err_arg(field):
    """ Generate an error for a missing required argument. """
    return err("Missing required argument: [hl]%s[/hl]" % field)


if __name__ == '__main__':
    main(sys.argv)