hblock

#!/usr/bin/python3

import os
import sys
import pickle
import shutil
import random
import re
import datetime
from os import system

home = os.path.expanduser('~')
usr = os.path.expanduser('/usr')
etc = os.path.expanduser('/etc')
hp = os.path.join(etc, 'hosts')
base = os.path.join(home, '.config', 'blocker')
act = os.path.join(base, 'act.pkl')
defs = os.path.join(base, 'def.pkl')
pas = os.path.join(base, 'pas.pkl')
tim = os.path.join(base, 'tim.pkl')
stp = os.path.join(base, 'stp.pkl')
wp = os.path.join(etc, 'hBlock', 'block-keywords')
root_base = os.path.join(etc, 'hBlock')
root_act = os.path.join(root_base, 'act.pkl')
debug = False

reserved = {'profile', '*', 'all'}
allowed = {'-', '_', ' '}
for i in range(65, 91):
    allowed.add(chr(i))
    allowed.add(chr(i).lower())
for i in range(9):
    allowed.add(str(i))

def import_list(fp):
    """
    Import a plaintext document as a list of strings where each item is a line
    of the document.

    Parameters
    ----------
    fp : pathlike
        File path of the document to read.

    Returns
    -------
    list
        List of strings of the lines of the document.
    """
    f = open(fp, 'r')
    l = f.read().split("\n")
    f.close()
    return l


words = import_list(wp)


def get_words(n=1):
    """
    Provides a string of random words.

    Parameters
    ----------
    n : int, optional
        Number of words to list, by default 1.

    Returns
    -------
    str
        String of the words separated by spaces.
    """
    string = ''
    for i in range(n):
        j = random.randint(0, 9999)
        string += words[j] + ' '
    string = string[:-1]
    return string


def check_name(name):
    """
    Checks if the provided block list name is valid.

    Parameters
    ----------
    name : str
        Name of the block list.

    Returns
    -------
    bool
        Whether or not the name is valid.
    """
    if name in reserved:
        print('Name ' + name + ' is reserved. Please try something else.')
        return False
    for c in name:
        if c not in allowed:
            print(
                'Name ' +
                name +
                ' uses invalid characters. Please try something else.')
            return False
    return True


def get_unlock():
    """
    Get the time that block lists will unlock.

    Returns
    -------
    str
        The unlock time.
    """
    try:
        load(tim)
    except BaseException:
        save_time()
    return load(tim)


def get_unpause():
    """
    Get thetime that block lists will unpause.

    Returns
    -------
    str
        The unpause time.
    """
    try:
        load(stp)
    except BaseException:
        save_pause()
    return load(stp)


def save_time(time=None):
    """
    Save a time for block lists to unlock.

    Parameters
    ----------
    time : str, optional
        The intended time, by default None. If left as None, will default to
        the current time.
    """
    if time is None:
        time = datetime.datetime.now()
    save(time, tim)


def save_pause(time=None):
    """
    Save a time for block lists to unpause.

    Parameters
    ----------
    time : str, optional
        The intended time, by default None. If left as None, will default to
        the current time.
    """
    if time is None:
        time = datetime.datetime.now()
    save(time, stp)


def check_time(message=True, time=None):
    """
    Check if the time is after the time for block lists to unlock.

    Parameters
    ----------
    message : bool, optional
        Locker passthrough.
    time : str, optional
        The intended time, by default None. If left as None, will default to
        the current time.

    Returns
    -------
    [type]
        [description]
    """
    unlock = get_unlock()
    if time is None:
        time = datetime.datetime.now()
    free = time > unlock
    if free:
        return free
    elif message:
        print('Cannot complete operation. Blocks are locked until ' +
              str(unlock) + '. Try again then.')
        return free
    else:
        return free


def plural(number, word):
    """
    Makes the word plural if need be.

    Parameters
    ----------
    number : int
        First half of output string.
    word : str
        Second half of output string.

    Returns
    -------
    str
        Number then word where the word is plural if the number isn't one.
    """
    if number == 1:
        return str(number) + ' ' + word
    else:
        return str(number) + ' ' + word + 's'


def root_refresh():
    """
    Refresh the root blocks from the local blocks.
    """
    fps = os.listdir(base)
    for p in fps:
        fp = os.path.join(base, p)
        rp = os.path.join(root_base, p)
        system('sudo cp ' + fp + ' ' + rp)


def check_block(block_dict):
    activity = get_activity()
    okay = True
    for block in block_dict.keys():
        try:
            if activity[block] and not block_dict[block]:
                okay = False
        except BaseException:
            pass
    return okay


def save(object, fp):
    pickle.dump(object, open(fp, 'wb'))
    root_path = os.path.join(root_base, os.path.basename(fp))
    system('sudo cp ' + fp + ' ' + root_path)


def load(fp):
    return pickle.load(open(fp, 'rb'))


def export_block(name, fp):
    cp = get_cp(name)
    shutil.copy(cp, fp)


def import_block(name, fp):
    shutil.copy(fp, base)
    activites = get_activity()
    activites[name] = False
    save(activites, act)
    system('sudo cp ' + act + ' ' + root_act)
    update_block(name)


def export_all(fp):
    new_base = os.path.join(base, 'blocker')
    try:
        shutil.rmtree(new_base)
    except BaseException:
        pass
    shutil.copytree(base, new_base)
    zipped_profile = os.path.join(base, 'hBlock_profile')
    try:
        os.remove(zipped_profile + '.zip')
    except BaseException:
        pass
    shutil.make_archive(zipped_profile, 'zip', new_base)
    zipped_profile += '.zip'
    fp = os.path.join(fp, 'hBlock_profile.zip')
    shutil.move(zipped_profile, fp)
    shutil.rmtree(new_base)


def import_all(fp):
    shutil.unpack_archive(fp, base)
    zipped_profile = os.path.join(base, 'hBlock_profile.zip')
    try:
        os.remove(zipped_profile)
    except BaseException:
        pass


def get_activity():
    return load(act)


def get_default():
    if not os.path.exists(defs):
        default = dict()
        save(default, defs)
    else:
        default = load(defs)
    return default


def get_keyword():
    return load(pas)


def set_keyword(kw):
    if not os.path.exists(pas):
        save(kw, pas)
    else:
        if require_keyword('Are you sure you want to change your keyword? '):
            if kw == '':
                os.remove(pas)
            else:
                print('Unblock keyword removed.')
            print("Unblock keyword word changed to '" + kw + ".'")
            save(kw, pas)
        else:
            print("Unblock keyword word unchanged.")


def require_keyword(message='', false_message='Operation canceled.'):
    if not os.path.exists(pas):
        return True
    else:
        kw = get_keyword()
        if kw == '':
            return True
        elif kw == 'random':
            kw = get_words(3)
        try:
            kn = int(kw)
            kw = get_words(kn)
        except BaseException:
            pass
        kw_check = input(
            message +
            "Enter '" +
            kw +
            "' to continue. Enter anything else to cancel.\n")
        if kw == kw_check:
            return True
        else:
            if false_message is not None:
                print(false_message)
            return False


def save_default():
    activity = get_activity()
    if not os.path.exists(defs):
        save(activity, defs)
    else:
        default = get_default()
        for name in activity.keys():
            default[name] = activity[name]
        save(default, defs)


def load_default():
    activity = get_activity()
    default = get_default()
    could_not_set = []
    for name in activity.keys():
        try:
            set_active(name, default[name])
        except BaseException:
            could_not_set.append(name)
    if len(could_not_set) == 1:
        print(
            "Block list '" +
            could_not_set[0] +
            "' was not in defaults. Run 'hblock save' to add it now.")
    elif len(could_not_set) > 1:
        s = 'Block lists '
        for name in could_not_set[:-1]:
            s += "'" + name + ",' "
        s += "and '" + could_not_set[-1] + "'"
        s += " were not in defaults. Run 'hblock save' to add them now."
        print(s)


def reset_default():
    os.remove(defs)


def get_bp(name, root=False, act_path=None):
    new_base = base
    if root:
        new_base = os.path.dirname(act_path)
    if check_name(name):
        return os.path.join(new_base, name + '.block')


def get_cp(name, root=False, act_path=None):
    new_base = base
    if root:
        new_base = os.path.dirname(act_path)
    if check_name(name):
        return os.path.join(new_base, name + '.cont')


def check_dir(fp):
    if not os.path.exists(fp):
        os.makedirs(fp)


def check_base():
    check_dir(base)
    if not os.path.exists(act):
        activity = {'default': False}
        save(activity, act)
        system('sudo cp ' + act + ' ' + root_act)
        new_block('default', [])


def setup():
    check_base()


def reset():
    os.system('rm -rf ' + str(base))
    setup()


def new_string(name, domains):
    if check_name(name):
        header = '# </blocker=' + str(name) + '>'
        footer = '# <blocker=' + str(name) + '/>'
        block_string = header + '\n'
        for domain in domains:
            block_string += '127.0.0.1	' + domain + '\n'
            block_string += '127.0.0.1	www.' + domain + '\n'
        block_string += footer + '\n'
        return block_string


def new_block(name, domains):
    if check_name(name):
        block_string = new_string(name, domains)
        check_base()
        bp = get_bp(name)
        block_file = open(bp, 'w')
        block_file.write(block_string)
        block_file.close()
        cp = get_cp(name)
        save(set(domains), cp)
        activity = get_activity()
        if name not in activity.keys():
            activity[name] = False
        save(activity, act)
        system('sudo cp ' + act + ' ' + root_act)


def delete_block(name):
    if check_name(name):
        bp = get_bp(name)
        cp = get_cp(name)
        os.remove(bp)
        os.remove(cp)
        activity = get_activity()
        activity.pop(name)
        save(activity, act)
        system('sudo cp ' + act + ' ' + root_act)


def expand_domains(add_domains):
    activity = get_activity()
    cur_domains = set(add_domains)
    add_domains = set(add_domains)
    for domain in cur_domains:
        if domain in activity.keys():
            add_domains = add_domains - {domain}
            new_cp = get_cp(domain)
            domains = load(new_cp)
            add_domains = add_domains | domains
    return add_domains


def add_to_block(name, add_domains):
    if check_name(name):
        cp = get_cp(name)
        old_domains = load(cp)
        add_domains = expand_domains(add_domains)
        new_domains = old_domains | add_domains
        save(new_domains, cp)
        update_block(name)


def rem_from_block(name, rem_domains):
    if check_name(name):
        cp = get_cp(name)
        old_domains = load(cp)
        rem_domains = set(rem_domains)
        new_domains = old_domains - rem_domains
        save(new_domains, cp)
        update_block(name)


def clear_block(name):
    if check_name(name):
        cp = get_cp(name)
        save(set(), cp)
        update_block(name)


def see_block(name):
    if check_name(name):
        cp = get_cp(name)
        domains = load(cp)
        active = get_activity()
        if active[name]:
            status = 'Blocked'
        else:
            status = 'Unblocked'
        print('Domains in ' + str(name) + ' [' + status + ']:')
        for domain in domains:
            print(' - ' + str(domain))


def see_all():
    active = get_activity()
    for name in active.keys():
        see_block(name)


def rename_block(old_name, new_name):
    if check_name(old_name) and check_name(new_name):
        old_bp = get_bp(old_name)
        old_cp = get_cp(old_name)
        new_cp = get_cp(new_name)
        domains = load(old_cp)
        os.remove(old_bp)
        os.remove(old_cp)
        save(domains, new_cp)
        update_block(new_name)
        activity = get_activity()
        activity[new_name] = activity[old_name]
        activity.pop(old_name)
        save(activity, act)
        system('sudo cp ' + act + ' ' + root_act)


def update_block(name):
    if check_name(name):
        cp = get_cp(name)
        try:
            domains = load(cp)
        except BaseException:
            domains = set()
        new_block(name, domains)
        update_hosts()


def set_active(name, on):
    if check_name(name):
        activity = get_activity()
        activity[name] = on
        save(activity, act)
        system('sudo cp ' + act + ' ' + root_act)
        update_block(name)


def see_active(search='*'):
    search = search.replace('*', '.')
    r = re.compile(search)
    activity = get_activity()
    names = list(activity.keys())
    results = list(filter(r.match, names))
    print('Block lists:')
    for name in results:
        space = 40 - len(name)
        if activity[name]:
            blocking = 'Blocked'
        else:
            blocking = 'Unblocked'
        print(' - ' + name + space * '.' + blocking)


def enable():
    activity = get_activity()
    for name in list(activity.keys()):
        set_active(name, True)


def disable():
    activity = get_activity()
    for name in list(activity.keys()):
        set_active(name, False)


def pause():
    print('This feature is currently a work in progress.')


def update_hosts():
    if debug:
        os.system('sudo ' + home + '/.local/bin/hblock root-update ' + act + ' --debug')
    else:
        os.system('sudo ' + home + '/.local/bin/hblock root-update ' + act)


def root_updater(act_path):
    activity = load(act_path)
    header = '\n# </hblock>\n'
    footer = '# <hblock/>\n'
    hosts = open(hp, 'r')
    h_string = hosts.read()
    if h_string.find(header) > -1:
        start = h_string.find(header)
        end = h_string.find(footer) + len(footer)
        old_update = h_string[start:end]
        h_string = h_string.replace(old_update, '')
    h_string += header
    for name in list(activity.keys()):
        if activity[name]:
            bp = get_bp(name, True, act_path)
            b = open(bp, 'r')
            b_string = b.read()
            b.close()
            h_string += b_string
    h_string += footer
    hosts.close()
    hosts = open(hp, 'w')
    hosts.write(h_string)
    hosts.close()


def lock(m, h, d):
    delta = datetime.timedelta(minutes=m, hours=h, days=d)
    now = datetime.datetime.now()
    unlock_time = now + delta
    if check_time(time=unlock_time):
        save_time(unlock_time)


def when_unlock():
    if check_time(message=False):
        print('There is no time lock on hBlock.')
    else:
        unlock = get_unlock()
        print('Blocks to unlock at ' + str(unlock) + '.')
        print('The current time is ' + str(datetime.datetime.now()) + '.')


command = sys.argv[1:]


def commands():
    if command[0] == 'active' or command[0] == '-a':
        try:
            see_active(command[1])
        except BaseException:
            see_active()
    elif command[0] == 'add':
        name = command[1]
        add_domains = command[2:]
        add_to_block(name, add_domains)
        print('Added ' +
              plural(len(expand_domains(add_domains)), 'domain') +
              ' to ' +
              name +
              '.')
    elif command[0] == 'remove' or command[0] == '-r':
        if check_time():
            name = command[1]
            rem_domains = command[2:]
            if require_keyword('Are you sure you want to remove ' +
                               plural(len(expand_domains(rem_domains)), 'domain') +
                               ' from ' +
                               name +
                               '? '):
                rem_from_block(name, rem_domains)
                print(
                    'Removed ' +
                    plural(
                        len(rem_domains),
                        'domain') +
                    ' from ' +
                    name +
                    '.')
    elif command[0] == 'show' or command[0] == '-s':
        if command[1] == 'all' or command[1] == '*':
            see_all()
        else:
            for name in command[1:]:
                see_block(name)
    elif command[0] == 'everything' or command[0] == 'all' or command[0] == '-e':
        see_all()
    elif command[0] == 'block' or command[0] == '-b':
        name = command[1]
        to_enable = command[2]
        if to_enable == 'on':
            to_enable = True
            print('Blocking ' + name + '.')
            set_active(name, to_enable)
        elif check_time() and require_keyword('Are you sure you want to turn blocking of ' + name + ' ' + to_enable + '? '):
            to_enable = False
            print('Unblocking ' + name + '.')
            set_active(name, to_enable)
    elif command[0] == 'new' or command[0] == '-n':
        name = command[1]
        domains = expand_domains(command[2:])
        new_block(name, domains)
        print('Created ' + name + ' with ' +
              plural(len(expand_domains(domains)), 'domain') + ' in it.')
    elif command[0] == 'delete' or command[0] == '-d':
        if check_time():
            name = command[1]
            if require_keyword(
                    'Are you sure you want to delete ' + name + '? '):
                print('Deleting ' + name + '.')
                delete_block(name)
    elif command[0] == 'rename':
        if check_time():
            old_name = command[1]
            new_name = command[2]
            if require_keyword(
                'Are you sure you want to rename ' +
                old_name +
                ' to ' +
                new_name +
                    '? '):
                rename_block(old_name, new_name)
                print('Renamed ' + old_name + ' to ' + new_name + '.')
    elif command[0] == 'setup':
        print('Attempting setup.')
        setup()
        print('Setup complete.')
    elif command[0] == 'reset':
        if check_time:
            if require_keyword(
                    'Are you sure you want to reset your entire hBlock profile? '):
                print('Attempting reset.')
                reset()
                print('Reset complete.')
    elif command[0] == 'update' or command[0] == '-u':
        if check_time:
            print('Forcing update of hosts.')
            update_hosts()
            print('Update of hosts complete.')
    elif command[0] == 'enable':
        if check_time:
            enable()
            print('Enabled ' + plural(len(get_activity().keys()), 'block list') + '.')
    elif command[0] == 'disable':
        if check_time():
            if require_keyword('Are you sure you want to disable ' +
                               plural(len(get_activity().keys()), 'block list') + '? '):
                disable()
                print('Disabled ' +
                      plural(len(get_activity().keys()), 'block list') +
                      '.')
    elif command[0] == 'import' or command[0] == '-i':
        if check_time():
            name = command[1]
            fp = command[2]
            if name == 'profile':
                import_all(fp)
            else:
                import_block(name, fp)
            root_refresh()
            print('Imported ' + name + ' from ' + fp + '.')
    elif command[0] == 'export':
        name = command[1]
        fp = command[2]
        if name == 'profile':
            export_all(fp)
        else:
            export_block(name, fp)
        print('Exported ' + name + ' to ' + fp + '.')
    elif command[0] == 'lock':
        time = command[1:]
        if 'd' in time:
            d = int(time[time.index('d') - 1])
        else:
            d = 0
        if 'h' in time:
            h = int(time[time.index('h') - 1])
        else:
            h = 0
        if 'm' in time:
            m = int(time[time.index('m') - 1])
        else:
            m = 0
        lock(m, h, d)
        print('Locking until ' + str(get_unlock()) + '.')
    elif command[0] == 'pause':
        time = command[1:]
        if 'd' in time:
            d = int(time[time.index('d') - 1])
        else:
            d = 0
        if 'h' in time:
            h = int(time[time.index('h') - 1])
        else:
            h = 0
        if 'm' in time:
            m = int(time[time.index('m') - 1])
        else:
            m = 0
        pause(m, h, d)
        print('Pausing until ' + str(get_unpause()) + '.')
    elif command[0] == 'load' or command[0] == '-l':
        if check_block(get_default()) or require_keyword(
                'Are you sure you want to load your default profile? '):
            load_default()
            print('Loaded default profile.')
    elif command[0] == 'save':
        if check_time():
            if require_keyword(
                    'Are you sure you want to save your current profile as the default? '):
                save_default()
                print('Saved default profile.')
    elif command[0] == 'refresh':
        if check_time():
            if require_keyword(
                    'Are you sure you want to reset your default profile? '):
                reset_default()
                print('Reset default profile.')
    elif command[0] == 'time' or command[0] == '-t':
        when_unlock()
    elif command[0] == 'keyword':
        pw = command[1]
        set_keyword(pw)
    elif command[0] == 'root-update':
        fp = command[1]
        root_updater(fp)
    elif command[0] == 'root-direct':
        fp = root_act
        root_updater(root_act)
    elif command[0] == 'root-refresh':
        root_refresh()
    elif command[0] == 'help' or command[0] == '-h':
        print(
            'Available commands:\n' +
            ' - active..................................Shows which block lists are active.\n' +
            ' - add [block list] [domains]..............Adds domains to block list.\n' +
            ' - remove [block list] [domains]...........Removes domains from block list.\n' +
            ' - show [block list].......................Shows all the domains in the block list.\n' +
            ' - everything..............................Shows all domains in all block lists.\n' +
            ' - block [block list] [on/off].............Turns the block list on or off.\n' +
            ' - enable..................................Turns on all block lists.\n' +
            ' - disable.................................Turns off all block lists.\n' +
            ' - pause [number] [m/h/d]..................Pause blocking for specified time.\n'
            ' - new [block list] [domains]..............Makes a new block list with specified domains.\n' +
            ' - delete [block list].....................Deletes block list.\n' +
            ' - rename [block list] [new name]..........Renames block list to the new name.\n' +
            ' - save....................................Saves current activity as default profile.\n' +
            ' - load....................................Loads default profile to current activity.\n' +
            " - lock [number] [m/h/d]...................Locks current profile so it can't be changed for\n" +
            "                                           specified time.\n" +
            " - keyword.................................Sets unblock keyword. Use 'random' or an integer to\n" +
            "                                           have it be different every time.\n" +
            ' - time....................................See when hBlock is scheduled to unlock.\n' +
            ' - setup...................................Automatically sets up hblock on first install.\n' +
            ' - reset...................................Reset hblock to its default state.\n' +
            ' - update..................................Forces update of hosts file.\n' +
            " - import [block list] [path]..............Imports block list with specified name Use 'profile'\n" +
            "                                           to import everything.\n" +
            " - export [block list] [path]..............Exports block list to specified path. Use 'profile'\n" +
            "                                           to export everything.\n" +
            ' - help....................................Shows available commands.')
    else:
        print(
            "Invalid command. Try entering 'hblock help' to see available commands.\n" +
            "If you think you entered a valid command, try running with the '--debug' flag at the end.")


if command[-1] == '--debug':
    debug = True
    command = command[:-1]
    print('~~~ Running in debug mode. ~~~')
    commands()
else:
    try:
        commands()
    except BaseException:
        print(
            "Invalid command. Try entering 'hblock help' to see available commands.\n" +
            "If you think you entered a valid command, try running with the '--debug' flag at the end.")