<commit_before>"""Consulate CLI commands"""
# pragma: no cover
import argparse
import base64
import json
import sys
import os

from requests import exceptions

import consulate
from consulate import adapters
from consulate import utils

def connection_error():
    """Common exit routine when consulate can't connect to Consul"""
    sys.stderr.write('ERROR: Could not connect to consul\n')
    sys.exit(1)

KV_PARSERS = [
    ('backup', 'Backup to stdout or a JSON file', [
        [['-f', '--file'],
         {'help': 'JSON file to read instead of stdin',
          'nargs': '?'}]]),
    ('restore', 'Restore from stdin or a JSON file', [
        [['-f', '--file'],
         {'help': 'JSON file to read instead of stdin',
          'nargs': '?'}],
        [['-n', '--no-replace'],
         {'help': 'Do not replace existing entries',
          'action': 'store_true'}]]),
    ('ls', 'List all of the keys', [
        [['-l', '--long'],
         {'help': 'Long format',
          'action': 'store_true'}]]),
    ('mkdir', 'Create a folder', [
        [['path'],
         {'help': 'The path to create'}]]),
    ('get', 'Get a key from the database', [
        [['key'],
         {'help': 'The key to get'}]]),
    ('set', 'Set a key in the database', [
        [['key'], {'help': 'The key to set'}],
        [['value'], {'help': 'The value of the key'}]]),
    ('rm', 'Remove a key from the database', [
        [['key'], {'help': 'The key to remove'}],
        [['-r', '--recurse'],
         {'help': 'Delete all keys prefixed with the specified key',
          'action': 'store_true'}]])]

def add_kv_args(parser):
    """Add the kv command and arguments.

    :param argparse.Subparser parser: parser

    """
    kv_parser = parser.add_parser('kv', help='Key/Value Database Utilities')

    subparsers = kv_parser.add_subparsers(dest='action',
                                          title='Key/Value Database Utilities')

    for (name, help_text, arguments) in KV_PARSERS:
        parser = subparsers.add_parser(name, help=help_text)
        for (args, kwargs) in arguments:
            parser.add_argument(*args, **kwargs)

def add_register_args(parser):
    """Add the register command and arguments.

    :param argparse.Subparser parser: parser

    """
    # Service registration
    registerp = parser.add_parser('register',
                                  help='Register a service for this node')
    registerp.add_argument('name', help='The service name')
    registerp.add_argument('-a', '--address', default=None,
                           help='Specify an address')
    registerp.add_argument('-p', '--port', default=None, type=int, help='Specify a port')
    registerp.add_argument('-s', '--service-id', default=None,
                           help='Specify a service ID')
    registerp.add_argument('-t', '--tags', default=[],
                           help='Specify a comma delimited list of tags')
    rsparsers = registerp.add_subparsers(dest='ctype',
                                         title='Service Check Options')
    check = rsparsers.add_parser('check',
                                 help='Define an external script-based check')
    check.add_argument('interval', default=10, type=int,
                       help='How often to run the check script')
    check.add_argument('path', default=None,
                       help='Path to the script invoked by Consul')
    httpcheck = rsparsers.add_parser('httpcheck',
                                 help='Define an HTTP-based check')
    httpcheck.add_argument('interval', default=10, type=int,
                       help='How often to run the check script')
    httpcheck.add_argument('url', default=None,
                       help='HTTP URL to be polled by Consul')
    rsparsers.add_parser('no-check', help='Do not enable service monitoring')
    ttl = rsparsers.add_parser('ttl', help='Define a duration based TTL check')
    ttl.add_argument('duration', type=int, default=10,
                     help='TTL duration for a service with missing check data')

def add_run_once_args(parser):
    """Add the run_once command and arguments.

    :param argparse.Subparser parser: parser

    """
    run_oncep = parser.add_parser('run_once',
                                  help='Lock this command')
    run_oncep.add_argument('prefix', help='the name of the lock which will be held in Consul.')
    run_oncep.add_argument('operation', nargs=argparse.REMAINDER, help='The command to lock')
    run_oncep.add_argument('-i', '--interval', default=None,
                           help='Hold the lock for X seconds')

def add_deregister_args(parser):
    """Add the deregister command and arguments.

    :param argparse.Subparser parser: parser

    """
    # Service registration
    registerp = parser.add_parser('deregister',
                                  help='Deregister a service for this node')
    registerp.add_argument('service_id', help='The service registration id')

def parse_cli_args():
    """Create the argument parser and add the arguments"""
    parser = argparse.ArgumentParser(description='CLI utilities for Consul')

    parser.add_argument('--api-scheme',
                        default='http',
                        help='The scheme to use for connecting to Consul with')
    parser.add_argument('--api-host',
                        help='The consul host to connect on')
    parser.add_argument('--api-port',
                        default=8500,
                        help='The consul API port to connect to')
    parser.add_argument('--datacenter',
                        dest='dc',
                        default=None,
                        help='The datacenter to specify for the connection')
    parser.add_argument('--token', default=None, help='ACL token')

    sparser = parser.add_subparsers(title='Commands', dest='command')
    add_register_args(sparser)
    add_deregister_args(sparser)
    add_kv_args(sparser)
    add_run_once_args(sparser)
    return parser.parse_args()

def kv_backup(consul, args):
    """Backup the Consul KV database

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    handle = open(args.file, 'w') if args.file else sys.stdout
    try:
        handle.write(json.dumps(consul.kv.records()) + '\n')
    except exceptions.ConnectionError:
        connection_error()

def kv_delete(consul, args):
    """Remove a key from the Consulate database

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    try:
        del consul.kv[args.key]
    except exceptions.ConnectionError:
        connection_error()

def kv_get(consul, args):
    """Get the value of a key from the Consul database

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    try:
        sys.stdout.write("%s\n" % consul.kv.get(args.key))
    except exceptions.ConnectionError:
        connection_error()

def kv_ls(consul, args):
    """List out the keys from the Consul KV database

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    try:
        for key in consul.kv.keys():
            if args.long:
                keylen = 0
                if consul.kv[key]:
                    keylen = len(consul.kv[key])
                print('{0:>14} {1}'.format(keylen, key))
            else:
                print(key)
    except exceptions.ConnectionError:
        connection_error()

def kv_mkdir(consul, args):
    """Make a key based path/directory in the KV database

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    if not args.path[:-1] == '/':
        args.path += '/'
    try:
        consul.kv.set(args.path, None)
    except exceptions.ConnectionError:
        connection_error()

def kv_restore(consul, args):
    """Restore the Consul KV store

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    handle = open(args.file, 'r') if args.file else sys.stdin
    data = json.load(handle)
    for row in data:
        if isinstance(row, dict):
            # translate raw api export to internal representation
            if row['Value'] is not None:
                row['Value'] = base64.b64decode(row['Value'])
            row = [row['Key'], row['Flags'], row['Value']]
        # Here's an awesome thing to make things work
        if not utils.PYTHON3 and isinstance(row[2], unicode):
            row[2] = row[2].encode('utf-8')
        try:
            consul.kv.set_record(row[0], row[1], row[2], not args.no_replace)
        except exceptions.ConnectionError:
            connection_error()

def kv_rm(consul, args):
    """Remove a key from the Consulate database

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    try:
        consul.kv.delete(args.key, args.recurse)
    except exceptions.ConnectionError:
        connection_error()

def kv_set(consul, args):
    """Set a value of a key int the Consul database

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    try:
        consul.kv[args.key] = args.value
    except exceptions.ConnectionError:
        connection_error()

def register(consul, args):
    """Handle service registration.

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    check = args.path if args.ctype == 'check' else None
    httpcheck = args.url if args.ctype == 'httpcheck' else None
    interval = '%ss' % args.interval if args.ctype in ['check', 'httpcheck'] else None
    ttl = '%ss' % args.duration if args.ctype == 'ttl' else None
    tags = args.tags.split(',') if args.tags else None
    try:
        consul.agent.service.register(args.name, args.service_id, args.address,
                                      args.port, tags, check, interval, ttl, httpcheck)
    except exceptions.ConnectionError:
        connection_error()

def deregister(consul, args):
    """Handle service deregistration.

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    try:
        consul.agent.service.deregister(args.service_id)
    except exceptions.ConnectionError:
        connection_error()

# Mapping dict to simplify the code in main()
KV_ACTIONS = {
    'backup': kv_backup,
    'del': kv_delete,
    'get': kv_get,
    'ls': kv_ls,
    'mkdir': kv_mkdir,
    'restore': kv_restore,
    'rm': kv_rm,
    'set': kv_set}

def run_once(consul, args):
    """Ensure only one process can run a command at a time

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    try:
        import time
        import subprocess

        session = consul.session.create()
        if not consul.kv.acquire_lock(args.prefix, session):
            sys.stdout.write('Cannot obtain the required lock. Exiting\n')
            sys.exit(2)
        if args.interval:
            now = int(time.time())
            last_run = consul.kv.get("{0}_last_run".format(args.prefix))
            if str(last_run) not in ['null', 'None'] and \
                int(last_run) + int(args.interval) > now:
                sys.stdout.write("Last run happened fewer than "\
                    "{0} second ago. Exiting\n".format(args.interval))
                consul.kv.release_lock(args.prefix, session)
                consul.session.destroy(session)
                return
            consul.kv["{0}_last_run".format(args.prefix)] = now
        consul.kv.release_lock(args.prefix, session)
        consul.session.destroy(session)

        # Should the subprocess return an error code, release the lock
        try:
            error = False
            subprocess.check_output(args.operation, stderr=subprocess.STDOUT)
        # If the subprocess fails
        except subprocess.CalledProcessError as e:
            sys.stdout.write('"{0}" exited with return code "{1}" and output {2}\n'.format(
                args.operation, e.returncode, e.output))
            error = True
        # If the command doesn't exist
        except OSError as e:
            sys.stdout.write('"{0}" command does not exist\n'.format(args.operation))
            error = True
        # Otherwise
        except Exception as e:
            sys.stdout.write('"{0}" exited with error "{1}"\n'.format(
                args.operation, e))
            error = True
        if error:
            sys.exit(1)
    except exceptions.ConnectionError:
        connection_error()

def main():
    """Entrypoint for the consulate cli application"""
    args = parse_cli_args()

    if args.api_scheme == 'http+unix':
        adapter = adapters.UnixSocketRequest
        port = None

        api_host = os.environ.get('CONSUL_HTTP_ADDR').replace('unix://', '')
        if args.api_host:
            api_host = args.api_host
    else:
        adapter = None
        port = args.api_port

        api_host = 'localhost'
        if args.api_host:
            api_host = args.api_host

    consul = consulate.Consul(api_host, port, args.dc,
                              args.token, args.api_scheme, adapter)

    if args.command == 'register':
        register(consul, args)
    elif args.command == 'deregister':
        deregister(consul, args)
    elif args.command == 'kv':
        KV_ACTIONS[args.action](consul, args)
    elif args.command == 'run_once':
        run_once(consul, args)

<commit_msg>Force port to be set to an int (#45)<commit_after>"""Consulate CLI commands"""
# pragma: no cover
import argparse
import base64
import json
import sys
import os

from requests import exceptions

import consulate
from consulate import adapters
from consulate import utils

def connection_error():
    """Common exit routine when consulate can't connect to Consul"""
    sys.stderr.write('ERROR: Could not connect to consul\n')
    sys.exit(1)

KV_PARSERS = [
    ('backup', 'Backup to stdout or a JSON file', [
        [['-f', '--file'],
         {'help': 'JSON file to read instead of stdin',
          'nargs': '?'}]]),
    ('restore', 'Restore from stdin or a JSON file', [
        [['-f', '--file'],
         {'help': 'JSON file to read instead of stdin',
          'nargs': '?'}],
        [['-n', '--no-replace'],
         {'help': 'Do not replace existing entries',
          'action': 'store_true'}]]),
    ('ls', 'List all of the keys', [
        [['-l', '--long'],
         {'help': 'Long format',
          'action': 'store_true'}]]),
    ('mkdir', 'Create a folder', [
        [['path'],
         {'help': 'The path to create'}]]),
    ('get', 'Get a key from the database', [
        [['key'],
         {'help': 'The key to get'}]]),
    ('set', 'Set a key in the database', [
        [['key'], {'help': 'The key to set'}],
        [['value'], {'help': 'The value of the key'}]]),
    ('rm', 'Remove a key from the database', [
        [['key'], {'help': 'The key to remove'}],
        [['-r', '--recurse'],
         {'help': 'Delete all keys prefixed with the specified key',
          'action': 'store_true'}]])]

def add_kv_args(parser):
    """Add the kv command and arguments.

    :param argparse.Subparser parser: parser

    """
    kv_parser = parser.add_parser('kv', help='Key/Value Database Utilities')

    subparsers = kv_parser.add_subparsers(dest='action',
                                          title='Key/Value Database Utilities')

    for (name, help_text, arguments) in KV_PARSERS:
        parser = subparsers.add_parser(name, help=help_text)
        for (args, kwargs) in arguments:
            parser.add_argument(*args, **kwargs)

def add_register_args(parser):
    """Add the register command and arguments.

    :param argparse.Subparser parser: parser

    """
    # Service registration
    registerp = parser.add_parser('register',
                                  help='Register a service for this node')
    registerp.add_argument('name', help='The service name')
    registerp.add_argument('-a', '--address', default=None,
                           help='Specify an address')
    registerp.add_argument('-p', '--port', default=None, type=int,
                           help='Specify a port')
    registerp.add_argument('-s', '--service-id', default=None,
                           help='Specify a service ID')
    registerp.add_argument('-t', '--tags', default=[],
                           help='Specify a comma delimited list of tags')
    rsparsers = registerp.add_subparsers(dest='ctype',
                                         title='Service Check Options')
    check = rsparsers.add_parser('check',
                                 help='Define an external script-based check')
    check.add_argument('interval', default=10, type=int,
                       help='How often to run the check script')
    check.add_argument('path', default=None,
                       help='Path to the script invoked by Consul')
    httpcheck = rsparsers.add_parser('httpcheck',
                                 help='Define an HTTP-based check')
    httpcheck.add_argument('interval', default=10, type=int,
                       help='How often to run the check script')
    httpcheck.add_argument('url', default=None,
                       help='HTTP URL to be polled by Consul')
    rsparsers.add_parser('no-check', help='Do not enable service monitoring')
    ttl = rsparsers.add_parser('ttl', help='Define a duration based TTL check')
    ttl.add_argument('duration', type=int, default=10,
                     help='TTL duration for a service with missing check data')

def add_run_once_args(parser):
    """Add the run_once command and arguments.

    :param argparse.Subparser parser: parser

    """
    run_oncep = parser.add_parser('run_once',
                                  help='Lock this command')
    run_oncep.add_argument('prefix', help='the name of the lock which will be held in Consul.')
    run_oncep.add_argument('operation', nargs=argparse.REMAINDER, help='The command to lock')
    run_oncep.add_argument('-i', '--interval', default=None,
                           help='Hold the lock for X seconds')

def add_deregister_args(parser):
    """Add the deregister command and arguments.

    :param argparse.Subparser parser: parser

    """
    # Service registration
    registerp = parser.add_parser('deregister',
                                  help='Deregister a service for this node')
    registerp.add_argument('service_id', help='The service registration id')

def parse_cli_args():
    """Create the argument parser and add the arguments"""
    parser = argparse.ArgumentParser(description='CLI utilities for Consul')

    parser.add_argument('--api-scheme',
                        default='http',
                        help='The scheme to use for connecting to Consul with')
    parser.add_argument('--api-host',
                        help='The consul host to connect on')
    parser.add_argument('--api-port',
                        default=8500,
                        help='The consul API port to connect to')
    parser.add_argument('--datacenter',
                        dest='dc',
                        default=None,
                        help='The datacenter to specify for the connection')
    parser.add_argument('--token', default=None, help='ACL token')

    sparser = parser.add_subparsers(title='Commands', dest='command')
    add_register_args(sparser)
    add_deregister_args(sparser)
    add_kv_args(sparser)
    add_run_once_args(sparser)
    return parser.parse_args()

def kv_backup(consul, args):
    """Backup the Consul KV database

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    handle = open(args.file, 'w') if args.file else sys.stdout
    try:
        handle.write(json.dumps(consul.kv.records()) + '\n')
    except exceptions.ConnectionError:
        connection_error()

def kv_delete(consul, args):
    """Remove a key from the Consulate database

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    try:
        del consul.kv[args.key]
    except exceptions.ConnectionError:
        connection_error()

def kv_get(consul, args):
    """Get the value of a key from the Consul database

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    try:
        sys.stdout.write("%s\n" % consul.kv.get(args.key))
    except exceptions.ConnectionError:
        connection_error()

def kv_ls(consul, args):
    """List out the keys from the Consul KV database

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    try:
        for key in consul.kv.keys():
            if args.long:
                keylen = 0
                if consul.kv[key]:
                    keylen = len(consul.kv[key])
                print('{0:>14} {1}'.format(keylen, key))
            else:
                print(key)
    except exceptions.ConnectionError:
        connection_error()

def kv_mkdir(consul, args):
    """Make a key based path/directory in the KV database

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    if not args.path[:-1] == '/':
        args.path += '/'
    try:
        consul.kv.set(args.path, None)
    except exceptions.ConnectionError:
        connection_error()

def kv_restore(consul, args):
    """Restore the Consul KV store

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    handle = open(args.file, 'r') if args.file else sys.stdin
    data = json.load(handle)
    for row in data:
        if isinstance(row, dict):
            # translate raw api export to internal representation
            if row['Value'] is not None:
                row['Value'] = base64.b64decode(row['Value'])
            row = [row['Key'], row['Flags'], row['Value']]
        # Here's an awesome thing to make things work
        if not utils.PYTHON3 and isinstance(row[2], unicode):
            row[2] = row[2].encode('utf-8')
        try:
            consul.kv.set_record(row[0], row[1], row[2], not args.no_replace)
        except exceptions.ConnectionError:
            connection_error()

def kv_rm(consul, args):
    """Remove a key from the Consulate database

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    try:
        consul.kv.delete(args.key, args.recurse)
    except exceptions.ConnectionError:
        connection_error()

def kv_set(consul, args):
    """Set a value of a key int the Consul database

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    try:
        consul.kv[args.key] = args.value
    except exceptions.ConnectionError:
        connection_error()

def register(consul, args):
    """Handle service registration.

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    check = args.path if args.ctype == 'check' else None
    httpcheck = args.url if args.ctype == 'httpcheck' else None
    interval = '%ss' % args.interval if args.ctype in ['check', 'httpcheck'] else None
    ttl = '%ss' % args.duration if args.ctype == 'ttl' else None
    tags = args.tags.split(',') if args.tags else None
    try:
        consul.agent.service.register(args.name, args.service_id, args.address,
                                      int(args.port), tags, check, interval,
                                      ttl, httpcheck)
    except exceptions.ConnectionError:
        connection_error()

def deregister(consul, args):
    """Handle service deregistration.

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    try:
        consul.agent.service.deregister(args.service_id)
    except exceptions.ConnectionError:
        connection_error()

# Mapping dict to simplify the code in main()
KV_ACTIONS = {
    'backup': kv_backup,
    'del': kv_delete,
    'get': kv_get,
    'ls': kv_ls,
    'mkdir': kv_mkdir,
    'restore': kv_restore,
    'rm': kv_rm,
    'set': kv_set}

def run_once(consul, args):
    """Ensure only one process can run a command at a time

    :param consulate.api_old.Consul consul: The Consul instance
    :param argparser.namespace args: The cli args

    """
    try:
        import time
        import subprocess

        session = consul.session.create()
        if not consul.kv.acquire_lock(args.prefix, session):
            sys.stdout.write('Cannot obtain the required lock. Exiting\n')
            sys.exit(2)
        if args.interval:
            now = int(time.time())
            last_run = consul.kv.get("{0}_last_run".format(args.prefix))
            if str(last_run) not in ['null', 'None'] and \
                int(last_run) + int(args.interval) > now:
                sys.stdout.write("Last run happened fewer than "\
                    "{0} second ago. Exiting\n".format(args.interval))
                consul.kv.release_lock(args.prefix, session)
                consul.session.destroy(session)
                return
            consul.kv["{0}_last_run".format(args.prefix)] = now
        consul.kv.release_lock(args.prefix, session)
        consul.session.destroy(session)

        # Should the subprocess return an error code, release the lock
        try:
            error = False
            subprocess.check_output(args.operation, stderr=subprocess.STDOUT)
        # If the subprocess fails
        except subprocess.CalledProcessError as e:
            sys.stdout.write('"{0}" exited with return code "{1}" and output {2}\n'.format(
                args.operation, e.returncode, e.output))
            error = True
        # If the command doesn't exist
        except OSError as e:
            sys.stdout.write('"{0}" command does not exist\n'.format(args.operation))
            error = True
        # Otherwise
        except Exception as e:
            sys.stdout.write('"{0}" exited with error "{1}"\n'.format(
                args.operation, e))
            error = True
        if error:
            sys.exit(1)
    except exceptions.ConnectionError:
        connection_error()

def main():
    """Entrypoint for the consulate cli application"""
    args = parse_cli_args()

    if args.api_scheme == 'http+unix':
        adapter = adapters.UnixSocketRequest
        port = None

        api_host = os.environ.get('CONSUL_HTTP_ADDR').replace('unix://', '')
        if args.api_host:
            api_host = args.api_host
    else:
        adapter = None
        port = args.api_port

        api_host = 'localhost'
        if args.api_host:
            api_host = args.api_host

    consul = consulate.Consul(api_host, port, args.dc,
                              args.token, args.api_scheme, adapter)

    if args.command == 'register':
        register(consul, args)
    elif args.command == 'deregister':
        deregister(consul, args)
    elif args.command == 'kv':
        KV_ACTIONS[args.action](consul, args)
    elif args.command == 'run_once':
        run_once(consul, args)

<|endoftext|>


<commit_before>from searx import settings, autocomplete
from searx.languages import language_codes as languages

COOKIE_MAX_AGE = 60 * 60 * 24 * 365 * 5  # 5 years
LANGUAGE_CODES = [l[0] for l in languages]
LANGUAGE_CODES.append('all')
DISABLED = 0
ENABLED = 1

class MissingArgumentException(Exception):
    pass

class ValidationException(Exception):
    pass

class Setting(object):
    """Base class of user settings"""

    def __init__(self, default_value, **kwargs):
        super(Setting, self).__init__()
        self.value = default_value
        for key, value in kwargs.iteritems():
            setattr(self, key, value)

        self._post_init()

    def _post_init(self):
        pass

    def parse(self, data):
        self.value = data

    def get_value(self):
        return self.value

    def save(self, name, resp):
        resp.set_cookie(name, bytes(self.value), max_age=COOKIE_MAX_AGE)

class StringSetting(Setting):
    """Setting of plain string values"""
    pass

class EnumStringSetting(Setting):
    """Setting of a value which can only come from the given choices"""

    def _post_init(self):
        if not hasattr(self, 'choices'):
            raise MissingArgumentException('Missing argument: choices')

        if self.value != '' and self.value not in self.choices:
            raise ValidationException('Invalid default value: {0}'.format(self.value))

    def parse(self, data):
        if data not in self.choices and data != self.value:
            raise ValidationException('Invalid choice: {0}'.format(data))
        self.value = data

class MultipleChoiceSetting(EnumStringSetting):
    """Setting of values which can only come from the given choices"""

    def _post_init(self):
        if not hasattr(self, 'choices'):
            raise MissingArgumentException('Missing argument: choices')
        for item in self.value:
            if item not in self.choices:
                raise ValidationException('Invalid default value: {0}'.format(self.value))

    def parse(self, data):
        if data == '':
            self.value = []
            return

        elements = data.split(',')
        for item in elements:
            if item not in self.choices:
                raise ValidationException('Invalid choice: {0}'.format(item))
        self.value = elements

    def parse_form(self, data):
        self.value = []
        for choice in data:
            if choice in self.choices and choice not in self.value:
                self.value.append(choice)

    def save(self, name, resp):
        resp.set_cookie(name, ','.join(self.value), max_age=COOKIE_MAX_AGE)

class MapSetting(Setting):
    """Setting of a value that has to be translated in order to be storable"""

    def _post_init(self):
        if not hasattr(self, 'map'):
            raise MissingArgumentException('missing argument: map')
        if self.value not in self.map.values():
            raise ValidationException('Invalid default value')

    def parse(self, data):
        if data not in self.map:
            raise ValidationException('Invalid choice: {0}'.format(data))
        self.value = self.map[data]
        self.key = data

    def save(self, name, resp):
        resp.set_cookie(name, bytes(self.key), max_age=COOKIE_MAX_AGE)

class SwitchableSetting(Setting):
    """ Base class for settings that can be turned on && off"""

    def _post_init(self):
        self.disabled = set()
        self.enabled = set()
        if not hasattr(self, 'choices'):
            raise MissingArgumentException('missing argument: choices')

    def transform_form_items(self, items):
        return items

    def transform_values(self, values):
        return values

    def parse_cookie(self, data):
        if data[DISABLED] != '':
            self.disabled = set(data[DISABLED].split(','))
        if data[ENABLED] != '':
            self.enabled = set(data[ENABLED].split(','))

    def parse_form(self, items):
        items = self.transform_form_items(items)

        self.disabled = set()
        self.enabled = set()
        for choice in self.choices:
            if choice['default_on']:
                if choice['id'] in items:
                    self.disabled.add(choice['id'])
            else:
                if choice['id'] not in items:
                    self.enabled.add(choice['id'])

    def save(self, resp):
        resp.set_cookie('disabled_{0}'.format(self.value), ','.join(self.disabled), max_age=COOKIE_MAX_AGE)
        resp.set_cookie('enabled_{0}'.format(self.value), ','.join(self.enabled), max_age=COOKIE_MAX_AGE)

    def get_disabled(self):
        disabled = self.disabled
        for choice in self.choices:
            if not choice['default_on'] and choice['id'] not in self.enabled:
                disabled.add(choice['id'])
        return self.transform_values(disabled)

    def get_enabled(self):
        enabled = self.enabled
        for choice in self.choices:
            if choice['default_on'] and choice['id'] not in self.disabled:
                enabled.add(choice['id'])
        return self.transform_values(enabled)

class EnginesSetting(SwitchableSetting):

    def _post_init(self):
        super(EnginesSetting, self)._post_init()
        transformed_choices = []
        for engine_name, engine in self.choices.iteritems():
            for category in engine.categories:
                transformed_choice = dict()
                transformed_choice['default_on'] = not engine.disabled
                transformed_choice['id'] = '{}__{}'.format(engine_name, category)
                transformed_choices.append(transformed_choice)
        self.choices = transformed_choices

    def transform_form_items(self, items):
        return [item[len('engine_'):].replace('_', ' ').replace('  ', '__') for item in items]

    def transform_values(self, values):
        if len(values) == 1 and next(iter(values)) == '':
            return list()
        transformed_values = []
        for value in values:
            engine, category = value.split('__')
            transformed_values.append((engine, category))
        return transformed_values

class PluginsSetting(SwitchableSetting):

    def _post_init(self):
        super(PluginsSetting, self)._post_init()
        transformed_choices = []
        for plugin in self.choices:
            transformed_choice = dict()
            transformed_choice['default_on'] = plugin.default_on
            transformed_choice['id'] = plugin.id
            transformed_choices.append(transformed_choice)
        self.choices = transformed_choices

    def transform_form_items(self, items):
        return [item[len('plugin_'):] for item in items]

class Preferences(object):
    """Stores, validates and saves preferences to cookies"""

    def __init__(self, themes, categories, engines, plugins):
        super(Preferences, self).__init__()

        self.key_value_settings = {'categories': MultipleChoiceSetting(['general'], choices=categories),
                                   'language': EnumStringSetting('all', choices=LANGUAGE_CODES),
                                   'locale': EnumStringSetting(settings['ui']['default_locale'],
                                                               choices=settings['locales'].keys()),
                                   'autocomplete': EnumStringSetting(settings['search']['autocomplete'],
                                                                     choices=autocomplete.backends.keys() + ['']),
                                   'image_proxy': MapSetting(settings['server']['image_proxy'],
                                                             map={'': settings['server']['image_proxy'],
                                                                  '0': False,
                                                                  '1': True}),
                                   'method': EnumStringSetting('POST', choices=('GET', 'POST')),
                                   'safesearch': MapSetting(settings['search']['safe_search'], map={'0': 0,
                                                                                '1': 1,
                                                                                '2': 2}),
                                   'theme': EnumStringSetting(settings['ui']['default_theme'], choices=themes),
                                   'results_on_new_tab': MapSetting(False, map={'0': False, '1': True})}

        self.engines = EnginesSetting('engines', choices=engines)
        self.plugins = PluginsSetting('plugins', choices=plugins)
        self.unknown_params = {}

    def parse_cookies(self, input_data):
        for user_setting_name, user_setting in input_data.iteritems():
            if user_setting_name in self.key_value_settings:
                self.key_value_settings[user_setting_name].parse(user_setting)
            elif user_setting_name == 'disabled_engines':
                self.engines.parse_cookie((input_data.get('disabled_engines', ''),
                                           input_data.get('enabled_engines', '')))
            elif user_setting_name == 'disabled_plugins':
                self.plugins.parse_cookie((input_data.get('disabled_plugins', ''),
                                           input_data.get('enabled_plugins', '')))

    def parse_form(self, input_data):
        disabled_engines = []
        enabled_categories = []
        disabled_plugins = []
        for user_setting_name, user_setting in input_data.iteritems():
            if user_setting_name in self.key_value_settings:
                self.key_value_settings[user_setting_name].parse(user_setting)
            elif user_setting_name.startswith('engine_'):
                disabled_engines.append(user_setting_name)
            elif user_setting_name.startswith('category_'):
                enabled_categories.append(user_setting_name[len('category_'):])
            elif user_setting_name.startswith('plugin_'):
                disabled_plugins.append(user_setting_name)
            else:
                self.unknown_params[user_setting_name] = user_setting
        self.key_value_settings['categories'].parse_form(enabled_categories)
        self.engines.parse_form(disabled_engines)
        self.plugins.parse_form(disabled_plugins)

    # cannot be used in case of engines or plugins
    def get_value(self, user_setting_name):
        if user_setting_name in self.key_value_settings:
            return self.key_value_settings[user_setting_name].get_value()

    def save(self, resp):
        for user_setting_name, user_setting in self.key_value_settings.iteritems():
            user_setting.save(user_setting_name, resp)
        self.engines.save(resp)
        self.plugins.save(resp)
        for k, v in self.unknown_params.items():
            resp.set_cookie(k, v, max_age=COOKIE_MAX_AGE)
        return resp
<commit_msg>preferences: refactor to check consistently input values<commit_after>from searx import settings, autocomplete
from searx.languages import language_codes as languages

COOKIE_MAX_AGE = 60 * 60 * 24 * 365 * 5  # 5 years
LANGUAGE_CODES = [l[0] for l in languages]
LANGUAGE_CODES.append('all')
DISABLED = 0
ENABLED = 1

class MissingArgumentException(Exception):
    pass

class ValidationException(Exception):
    pass

class Setting(object):
    """Base class of user settings"""

    def __init__(self, default_value, **kwargs):
        super(Setting, self).__init__()
        self.value = default_value
        for key, value in kwargs.iteritems():
            setattr(self, key, value)

        self._post_init()

    def _post_init(self):
        pass

    def parse(self, data):
        self.value = data

    def get_value(self):
        return self.value

    def save(self, name, resp):
        resp.set_cookie(name, bytes(self.value), max_age=COOKIE_MAX_AGE)

class StringSetting(Setting):
    """Setting of plain string values"""
    pass

class EnumStringSetting(Setting):
    """Setting of a value which can only come from the given choices"""

    def _validate_selection(self, selection):
        if selection not in self.choices:
            raise ValidationException('Invalid value: "{0}"'.format(selection))

    def _post_init(self):
        if not hasattr(self, 'choices'):
            raise MissingArgumentException('Missing argument: choices')
        self._validate_selection(self.value)

    def parse(self, data):
        self._validate_selection(data)
        self.value = data

class MultipleChoiceSetting(EnumStringSetting):
    """Setting of values which can only come from the given choices"""

    def _validate_selections(self, selections):
        for item in selections:
            if item not in self.choices:
                raise ValidationException('Invalid value: "{0}"'.format(selections))

    def _post_init(self):
        if not hasattr(self, 'choices'):
            raise MissingArgumentException('Missing argument: choices')
        self._validate_selections(self.value)

    def parse(self, data):
        if data == '':
            self.value = []
            return

        elements = data.split(',')
        self._validate_selections(elements)
        self.value = elements

    def parse_form(self, data):
        self.value = []
        for choice in data:
            if choice in self.choices and choice not in self.value:
                self.value.append(choice)

    def save(self, name, resp):
        resp.set_cookie(name, ','.join(self.value), max_age=COOKIE_MAX_AGE)

class MapSetting(Setting):
    """Setting of a value that has to be translated in order to be storable"""

    def _post_init(self):
        if not hasattr(self, 'map'):
            raise MissingArgumentException('missing argument: map')
        if self.value not in self.map.values():
            raise ValidationException('Invalid default value')

    def parse(self, data):
        if data not in self.map:
            raise ValidationException('Invalid choice: {0}'.format(data))
        self.value = self.map[data]
        self.key = data

    def save(self, name, resp):
        resp.set_cookie(name, bytes(self.key), max_age=COOKIE_MAX_AGE)

class SwitchableSetting(Setting):
    """ Base class for settings that can be turned on && off"""

    def _post_init(self):
        self.disabled = set()
        self.enabled = set()
        if not hasattr(self, 'choices'):
            raise MissingArgumentException('missing argument: choices')

    def transform_form_items(self, items):
        return items

    def transform_values(self, values):
        return values

    def parse_cookie(self, data):
        if data[DISABLED] != '':
            self.disabled = set(data[DISABLED].split(','))
        if data[ENABLED] != '':
            self.enabled = set(data[ENABLED].split(','))

    def parse_form(self, items):
        items = self.transform_form_items(items)

        self.disabled = set()
        self.enabled = set()
        for choice in self.choices:
            if choice['default_on']:
                if choice['id'] in items:
                    self.disabled.add(choice['id'])
            else:
                if choice['id'] not in items:
                    self.enabled.add(choice['id'])

    def save(self, resp):
        resp.set_cookie('disabled_{0}'.format(self.value), ','.join(self.disabled), max_age=COOKIE_MAX_AGE)
        resp.set_cookie('enabled_{0}'.format(self.value), ','.join(self.enabled), max_age=COOKIE_MAX_AGE)

    def get_disabled(self):
        disabled = self.disabled
        for choice in self.choices:
            if not choice['default_on'] and choice['id'] not in self.enabled:
                disabled.add(choice['id'])
        return self.transform_values(disabled)

    def get_enabled(self):
        enabled = self.enabled
        for choice in self.choices:
            if choice['default_on'] and choice['id'] not in self.disabled:
                enabled.add(choice['id'])
        return self.transform_values(enabled)

class EnginesSetting(SwitchableSetting):

    def _post_init(self):
        super(EnginesSetting, self)._post_init()
        transformed_choices = []
        for engine_name, engine in self.choices.iteritems():
            for category in engine.categories:
                transformed_choice = dict()
                transformed_choice['default_on'] = not engine.disabled
                transformed_choice['id'] = '{}__{}'.format(engine_name, category)
                transformed_choices.append(transformed_choice)
        self.choices = transformed_choices

    def transform_form_items(self, items):
        return [item[len('engine_'):].replace('_', ' ').replace('  ', '__') for item in items]

    def transform_values(self, values):
        if len(values) == 1 and next(iter(values)) == '':
            return list()
        transformed_values = []
        for value in values:
            engine, category = value.split('__')
            transformed_values.append((engine, category))
        return transformed_values

class PluginsSetting(SwitchableSetting):

    def _post_init(self):
        super(PluginsSetting, self)._post_init()
        transformed_choices = []
        for plugin in self.choices:
            transformed_choice = dict()
            transformed_choice['default_on'] = plugin.default_on
            transformed_choice['id'] = plugin.id
            transformed_choices.append(transformed_choice)
        self.choices = transformed_choices

    def transform_form_items(self, items):
        return [item[len('plugin_'):] for item in items]

class Preferences(object):
    """Stores, validates and saves preferences to cookies"""

    def __init__(self, themes, categories, engines, plugins):
        super(Preferences, self).__init__()

        self.key_value_settings = {'categories': MultipleChoiceSetting(['general'], choices=categories),
                                   'language': EnumStringSetting('all', choices=LANGUAGE_CODES),
                                   'locale': EnumStringSetting(settings['ui']['default_locale'],
                                                               choices=settings['locales'].keys() + ['']),
                                   'autocomplete': EnumStringSetting(settings['search']['autocomplete'],
                                                                     choices=autocomplete.backends.keys() + ['']),
                                   'image_proxy': MapSetting(settings['server']['image_proxy'],
                                                             map={'': settings['server']['image_proxy'],
                                                                  '0': False,
                                                                  '1': True}),
                                   'method': EnumStringSetting('POST', choices=('GET', 'POST')),
                                   'safesearch': MapSetting(settings['search']['safe_search'], map={'0': 0,
                                                                                '1': 1,
                                                                                '2': 2}),
                                   'theme': EnumStringSetting(settings['ui']['default_theme'], choices=themes),
                                   'results_on_new_tab': MapSetting(False, map={'0': False, '1': True})}

        self.engines = EnginesSetting('engines', choices=engines)
        self.plugins = PluginsSetting('plugins', choices=plugins)
        self.unknown_params = {}

    def parse_cookies(self, input_data):
        for user_setting_name, user_setting in input_data.iteritems():
            if user_setting_name in self.key_value_settings:
                self.key_value_settings[user_setting_name].parse(user_setting)
            elif user_setting_name == 'disabled_engines':
                self.engines.parse_cookie((input_data.get('disabled_engines', ''),
                                           input_data.get('enabled_engines', '')))
            elif user_setting_name == 'disabled_plugins':
                self.plugins.parse_cookie((input_data.get('disabled_plugins', ''),
                                           input_data.get('enabled_plugins', '')))

    def parse_form(self, input_data):
        disabled_engines = []
        enabled_categories = []
        disabled_plugins = []
        for user_setting_name, user_setting in input_data.iteritems():
            if user_setting_name in self.key_value_settings:
                self.key_value_settings[user_setting_name].parse(user_setting)
            elif user_setting_name.startswith('engine_'):
                disabled_engines.append(user_setting_name)
            elif user_setting_name.startswith('category_'):
                enabled_categories.append(user_setting_name[len('category_'):])
            elif user_setting_name.startswith('plugin_'):
                disabled_plugins.append(user_setting_name)
            else:
                self.unknown_params[user_setting_name] = user_setting
        self.key_value_settings['categories'].parse_form(enabled_categories)
        self.engines.parse_form(disabled_engines)
        self.plugins.parse_form(disabled_plugins)

    # cannot be used in case of engines or plugins
    def get_value(self, user_setting_name):
        if user_setting_name in self.key_value_settings:
            return self.key_value_settings[user_setting_name].get_value()

    def save(self, resp):
        for user_setting_name, user_setting in self.key_value_settings.iteritems():
            user_setting.save(user_setting_name, resp)
        self.engines.save(resp)
        self.plugins.save(resp)
        for k, v in self.unknown_params.items():
            resp.set_cookie(k, v, max_age=COOKIE_MAX_AGE)
        return resp
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import binascii
import datetime
import base64
import os
import scrypt
import pyotp
import qrcode
# Using svg because it doesn't require additional dependencies
import qrcode.image.svg

# Find the best implementation available on this platform
try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO  # type: ignore

from flask import current_app
from jinja2 import Markup
from sqlalchemy import ForeignKey
from sqlalchemy.orm import relationship, backref
from sqlalchemy import Column, Integer, String, Boolean, DateTime, Binary
from sqlalchemy.orm.exc import MultipleResultsFound, NoResultFound

from db import db

LOGIN_HARDENING = True
if os.environ.get('SECUREDROP_ENV') == 'test':
    LOGIN_HARDENING = False

def get_one_or_else(query, logger, failure_method):
    try:
        return query.one()
    except MultipleResultsFound as e:
        logger.error(
            "Found multiple while executing %s when one was expected: %s" %
            (query, e, ))
        failure_method(500)
    except NoResultFound as e:
        logger.error("Found none when one was expected: %s" % (e,))
        failure_method(404)

class Source(db.Model):
    __tablename__ = 'sources'
    id = Column(Integer, primary_key=True)
    filesystem_id = Column(String(96), unique=True)
    journalist_designation = Column(String(255), nullable=False)
    flagged = Column(Boolean, default=False)
    last_updated = Column(DateTime, default=datetime.datetime.utcnow)
    star = relationship("SourceStar", uselist=False, backref="source")

    # sources are "pending" and don't get displayed to journalists until they
    # submit something
    pending = Column(Boolean, default=True)

    # keep track of how many interactions have happened, for filenames
    interaction_count = Column(Integer, default=0, nullable=False)

    # Don't create or bother checking excessively long codenames to prevent DoS
    NUM_WORDS = 7
    MAX_CODENAME_LEN = 128

    def __init__(self, filesystem_id=None, journalist_designation=None):
        self.filesystem_id = filesystem_id
        self.journalist_designation = journalist_designation

    def __repr__(self):
        return '<Source %r>' % (self.journalist_designation)

    @property
    def journalist_filename(self):
        valid_chars = 'abcdefghijklmnopqrstuvwxyz1234567890-_'
        return ''.join([c for c in self.journalist_designation.lower().replace(
            ' ', '_') if c in valid_chars])

    def documents_messages_count(self):
        try:
            return self.docs_msgs_count
        except AttributeError:
            self.docs_msgs_count = {'messages': 0, 'documents': 0}
            for submission in self.submissions:
                if submission.filename.endswith('msg.gpg'):
                    self.docs_msgs_count['messages'] += 1
                elif (submission.filename.endswith('doc.gz.gpg') or
                      submission.filename.endswith('doc.zip.gpg')):
                    self.docs_msgs_count['documents'] += 1
            return self.docs_msgs_count

    @property
    def collection(self):
        """Return the list of submissions and replies for this source, sorted
        in ascending order by the filename/interaction count."""
        collection = []
        collection.extend(self.submissions)
        collection.extend(self.replies)
        collection.sort(key=lambda x: int(x.filename.split('-')[0]))
        return collection

class Submission(db.Model):
    __tablename__ = 'submissions'
    id = Column(Integer, primary_key=True)
    source_id = Column(Integer, ForeignKey('sources.id'))
    source = relationship(
        "Source",
        backref=backref("submissions", order_by=id, cascade="delete")
        )

    filename = Column(String(255), nullable=False)
    size = Column(Integer, nullable=False)
    downloaded = Column(Boolean, default=False)

    def __init__(self, source, filename):
        self.source_id = source.id
        self.filename = filename
        self.size = os.stat(current_app.storage.path(source.filesystem_id,
                                                     filename)).st_size

    def __repr__(self):
        return '<Submission %r>' % (self.filename)

class Reply(db.Model):
    __tablename__ = "replies"
    id = Column(Integer, primary_key=True)

    journalist_id = Column(Integer, ForeignKey('journalists.id'))
    journalist = relationship(
        "Journalist",
        backref=backref(
            'replies',
            order_by=id))

    source_id = Column(Integer, ForeignKey('sources.id'))
    source = relationship(
        "Source",
        backref=backref("replies", order_by=id, cascade="delete")
        )

    filename = Column(String(255), nullable=False)
    size = Column(Integer, nullable=False)

    def __init__(self, journalist, source, filename):
        self.journalist_id = journalist.id
        self.source_id = source.id
        self.filename = filename
        self.size = os.stat(current_app.storage.path(source.filesystem_id,
                                                     filename)).st_size

    def __repr__(self):
        return '<Reply %r>' % (self.filename)

class SourceStar(db.Model):
    __tablename__ = 'source_stars'
    id = Column("id", Integer, primary_key=True)
    source_id = Column("source_id", Integer, ForeignKey('sources.id'))
    starred = Column("starred", Boolean, default=True)

    def __eq__(self, other):
        if isinstance(other, SourceStar):
            return (self.source_id == other.source_id and
                    self.id == other.id and self.starred == other.starred)
        return NotImplemented

    def __init__(self, source, starred=True):
        self.source_id = source.id
        self.starred = starred

class InvalidUsernameException(Exception):

    """Raised when a user logs in with an invalid username"""

class LoginThrottledException(Exception):

    """Raised when a user attempts to log in
    too many times in a given time period"""

class WrongPasswordException(Exception):

    """Raised when a user logs in with an incorrect password"""

class BadTokenException(Exception):

    """Raised when a user logins in with an incorrect TOTP token"""

class PasswordError(Exception):

    """Generic error for passwords that are invalid.
    """

class InvalidPasswordLength(PasswordError):
    """Raised when attempting to create a Journalist or log in with an invalid
       password length.
    """

    def __init__(self, password):
        self.pw_len = len(password)

    def __str__(self):
        if self.pw_len > Journalist.MAX_PASSWORD_LEN:
            return "Password too long (len={})".format(self.pw_len)
        if self.pw_len < Journalist.MIN_PASSWORD_LEN:
            return "Password needs to be at least {} characters".format(
                Journalist.MIN_PASSWORD_LEN
            )

class NonDicewarePassword(PasswordError):

    """Raised when attempting to validate a password that is not diceware-like
    """

class Journalist(db.Model):
    __tablename__ = "journalists"
    id = Column(Integer, primary_key=True)
    username = Column(String(255), nullable=False, unique=True)
    pw_salt = Column(Binary(32))
    pw_hash = Column(Binary(256))
    is_admin = Column(Boolean)

    otp_secret = Column(String(16), default=pyotp.random_base32)
    is_totp = Column(Boolean, default=True)
    hotp_counter = Column(Integer, default=0)
    last_token = Column(String(6))

    created_on = Column(DateTime, default=datetime.datetime.utcnow)
    last_access = Column(DateTime)
    login_attempts = relationship(
        "JournalistLoginAttempt",
        backref="journalist")

    MIN_USERNAME_LEN = 3

    def __init__(self, username, password, is_admin=False, otp_secret=None):
        self.check_username_acceptable(username)
        self.username = username
        self.set_password(password)
        self.is_admin = is_admin

        if otp_secret:
            self.set_hotp_secret(otp_secret)

    def __repr__(self):
        return "<Journalist {0}{1}>".format(
            self.username,
            " [admin]" if self.is_admin else "")

    def _gen_salt(self, salt_bytes=32):
        return os.urandom(salt_bytes)

    _SCRYPT_PARAMS = dict(N=2**14, r=8, p=1)

    def _scrypt_hash(self, password, salt, params=None):
        if not params:
            params = self._SCRYPT_PARAMS
        return scrypt.hash(str(password), salt, **params)

    MAX_PASSWORD_LEN = 128
    MIN_PASSWORD_LEN = 14

    def set_password(self, password):
        self.check_password_acceptable(password)

        # Don't do anything if user's password hasn't changed.
        if self.pw_hash and self.valid_password(password):
            return

        self.pw_salt = self._gen_salt()
        self.pw_hash = self._scrypt_hash(password, self.pw_salt)

    @classmethod
    def check_username_acceptable(cls, username):
        if len(username) < cls.MIN_USERNAME_LEN:
            raise InvalidUsernameException(
                        'Username "{}" must be at least {} characters long.'
                        .format(username, cls.MIN_USERNAME_LEN))

    @classmethod
    def check_password_acceptable(cls, password):
        # Enforce a reasonable maximum length for passwords to avoid DoS
        if len(password) > cls.MAX_PASSWORD_LEN:
            raise InvalidPasswordLength(password)

        # Enforce a reasonable minimum length for new passwords
        if len(password) < cls.MIN_PASSWORD_LEN:
            raise InvalidPasswordLength(password)

        # Ensure all passwords are "diceware-like"
        if len(password.split()) < 7:
            raise NonDicewarePassword()

    def valid_password(self, password):
        # Avoid hashing passwords that are over the maximum length
        if len(password) > self.MAX_PASSWORD_LEN:
            raise InvalidPasswordLength(password)
        # No check on minimum password length here because some passwords
        # may have been set prior to setting the minimum password length.
        return pyotp.utils.compare_digest(
            self._scrypt_hash(password, self.pw_salt),
            self.pw_hash)

    def regenerate_totp_shared_secret(self):
        self.otp_secret = pyotp.random_base32()

    def set_hotp_secret(self, otp_secret):
        self.otp_secret = base64.b32encode(
            binascii.unhexlify(
                otp_secret.replace(
                    " ",
                    "")))
        self.is_totp = False
        self.hotp_counter = 0

    @property
    def totp(self):
        if self.is_totp:
            return pyotp.TOTP(self.otp_secret)
        else:
            raise ValueError('{} is not using TOTP'.format(self))

    @property
    def hotp(self):
        if not self.is_totp:
            return pyotp.HOTP(self.otp_secret)
        else:
            raise ValueError('{} is not using HOTP'.format(self))

    @property
    def shared_secret_qrcode(self):
        uri = self.totp.provisioning_uri(
            self.username,
            issuer_name="SecureDrop")

        qr = qrcode.QRCode(
            box_size=15,
            image_factory=qrcode.image.svg.SvgPathImage
        )
        qr.add_data(uri)
        img = qr.make_image()

        svg_out = StringIO()
        img.save(svg_out)
        return Markup(svg_out.getvalue())

    @property
    def formatted_otp_secret(self):
        """The OTP secret is easier to read and manually enter if it is all
        lowercase and split into four groups of four characters. The secret is
        base32-encoded, so it is case insensitive."""
        sec = self.otp_secret
        chunks = [sec[i:i + 4] for i in range(0, len(sec), 4)]
        return ' '.join(chunks).lower()

    def _format_token(self, token):
        """Strips from authentication tokens the whitespace
        that many clients add for readability"""
        return ''.join(token.split())

    def verify_token(self, token):
        token = self._format_token(token)

        # Store latest token to prevent OTP token reuse
        self.last_token = token
        db.session.commit()

        if self.is_totp:
            # Also check the given token against the previous and next
            # valid tokens, to compensate for potential time skew
            # between the client and the server. The total valid
            # window is 1:30s.
            return self.totp.verify(token, valid_window=1)
        else:
            for counter_val in range(
                    self.hotp_counter,
                    self.hotp_counter + 20):
                if self.hotp.verify(token, counter_val):
                    self.hotp_counter = counter_val + 1
                    db.session.commit()
                    return True
            return False

    _LOGIN_ATTEMPT_PERIOD = 60  # seconds
    _MAX_LOGIN_ATTEMPTS_PER_PERIOD = 5

    @classmethod
    def throttle_login(cls, user):
        # Record the login attempt...
        login_attempt = JournalistLoginAttempt(user)
        db.session.add(login_attempt)
        db.session.commit()

        # ...and reject it if they have exceeded the threshold
        login_attempt_period = datetime.datetime.utcnow() - \
            datetime.timedelta(seconds=cls._LOGIN_ATTEMPT_PERIOD)
        attempts_within_period = JournalistLoginAttempt.query.filter(
            JournalistLoginAttempt.journalist_id == user.id).filter(
            JournalistLoginAttempt.timestamp > login_attempt_period).all()
        if len(attempts_within_period) > cls._MAX_LOGIN_ATTEMPTS_PER_PERIOD:
            raise LoginThrottledException(
                "throttled ({} attempts in last {} seconds)".format(
                    len(attempts_within_period),
                    cls._LOGIN_ATTEMPT_PERIOD))

    @classmethod
    def login(cls, username, password, token):
        try:
            user = Journalist.query.filter_by(username=username).one()
        except NoResultFound:
            raise InvalidUsernameException(
                "invalid username '{}'".format(username))

        if LOGIN_HARDENING:
            cls.throttle_login(user)

            # Prevent TOTP token reuse
            if user.last_token is not None:
                if pyotp.utils.compare_digest(token, user.last_token):
                    raise BadTokenException("previously used token "
                                            "{}".format(token))
        if not user.verify_token(token):
            raise BadTokenException("invalid token")
        if not user.valid_password(password):
            raise WrongPasswordException("invalid password")
        return user

class JournalistLoginAttempt(db.Model):

    """This model keeps track of journalist's login attempts so we can
    rate limit them in order to prevent attackers from brute forcing
    passwords or two-factor tokens."""
    __tablename__ = "journalist_login_attempt"
    id = Column(Integer, primary_key=True)
    timestamp = Column(DateTime, default=datetime.datetime.utcnow)
    journalist_id = Column(Integer, ForeignKey('journalists.id'))

    def __init__(self, journalist):
        self.journalist_id = journalist.id
<commit_msg>source: do not set last_updated default to now<commit_after># -*- coding: utf-8 -*-
import binascii
import datetime
import base64
import os
import scrypt
import pyotp
import qrcode
# Using svg because it doesn't require additional dependencies
import qrcode.image.svg

# Find the best implementation available on this platform
try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO  # type: ignore

from flask import current_app
from jinja2 import Markup
from sqlalchemy import ForeignKey
from sqlalchemy.orm import relationship, backref
from sqlalchemy import Column, Integer, String, Boolean, DateTime, Binary
from sqlalchemy.orm.exc import MultipleResultsFound, NoResultFound

from db import db

LOGIN_HARDENING = True
if os.environ.get('SECUREDROP_ENV') == 'test':
    LOGIN_HARDENING = False

def get_one_or_else(query, logger, failure_method):
    try:
        return query.one()
    except MultipleResultsFound as e:
        logger.error(
            "Found multiple while executing %s when one was expected: %s" %
            (query, e, ))
        failure_method(500)
    except NoResultFound as e:
        logger.error("Found none when one was expected: %s" % (e,))
        failure_method(404)

class Source(db.Model):
    __tablename__ = 'sources'
    id = Column(Integer, primary_key=True)
    filesystem_id = Column(String(96), unique=True)
    journalist_designation = Column(String(255), nullable=False)
    flagged = Column(Boolean, default=False)
    last_updated = Column(DateTime)
    star = relationship("SourceStar", uselist=False, backref="source")

    # sources are "pending" and don't get displayed to journalists until they
    # submit something
    pending = Column(Boolean, default=True)

    # keep track of how many interactions have happened, for filenames
    interaction_count = Column(Integer, default=0, nullable=False)

    # Don't create or bother checking excessively long codenames to prevent DoS
    NUM_WORDS = 7
    MAX_CODENAME_LEN = 128

    def __init__(self, filesystem_id=None, journalist_designation=None):
        self.filesystem_id = filesystem_id
        self.journalist_designation = journalist_designation

    def __repr__(self):
        return '<Source %r>' % (self.journalist_designation)

    @property
    def journalist_filename(self):
        valid_chars = 'abcdefghijklmnopqrstuvwxyz1234567890-_'
        return ''.join([c for c in self.journalist_designation.lower().replace(
            ' ', '_') if c in valid_chars])

    def documents_messages_count(self):
        try:
            return self.docs_msgs_count
        except AttributeError:
            self.docs_msgs_count = {'messages': 0, 'documents': 0}
            for submission in self.submissions:
                if submission.filename.endswith('msg.gpg'):
                    self.docs_msgs_count['messages'] += 1
                elif (submission.filename.endswith('doc.gz.gpg') or
                      submission.filename.endswith('doc.zip.gpg')):
                    self.docs_msgs_count['documents'] += 1
            return self.docs_msgs_count

    @property
    def collection(self):
        """Return the list of submissions and replies for this source, sorted
        in ascending order by the filename/interaction count."""
        collection = []
        collection.extend(self.submissions)
        collection.extend(self.replies)
        collection.sort(key=lambda x: int(x.filename.split('-')[0]))
        return collection

class Submission(db.Model):
    __tablename__ = 'submissions'
    id = Column(Integer, primary_key=True)
    source_id = Column(Integer, ForeignKey('sources.id'))
    source = relationship(
        "Source",
        backref=backref("submissions", order_by=id, cascade="delete")
        )

    filename = Column(String(255), nullable=False)
    size = Column(Integer, nullable=False)
    downloaded = Column(Boolean, default=False)

    def __init__(self, source, filename):
        self.source_id = source.id
        self.filename = filename
        self.size = os.stat(current_app.storage.path(source.filesystem_id,
                                                     filename)).st_size

    def __repr__(self):
        return '<Submission %r>' % (self.filename)

class Reply(db.Model):
    __tablename__ = "replies"
    id = Column(Integer, primary_key=True)

    journalist_id = Column(Integer, ForeignKey('journalists.id'))
    journalist = relationship(
        "Journalist",
        backref=backref(
            'replies',
            order_by=id))

    source_id = Column(Integer, ForeignKey('sources.id'))
    source = relationship(
        "Source",
        backref=backref("replies", order_by=id, cascade="delete")
        )

    filename = Column(String(255), nullable=False)
    size = Column(Integer, nullable=False)

    def __init__(self, journalist, source, filename):
        self.journalist_id = journalist.id
        self.source_id = source.id
        self.filename = filename
        self.size = os.stat(current_app.storage.path(source.filesystem_id,
                                                     filename)).st_size

    def __repr__(self):
        return '<Reply %r>' % (self.filename)

class SourceStar(db.Model):
    __tablename__ = 'source_stars'
    id = Column("id", Integer, primary_key=True)
    source_id = Column("source_id", Integer, ForeignKey('sources.id'))
    starred = Column("starred", Boolean, default=True)

    def __eq__(self, other):
        if isinstance(other, SourceStar):
            return (self.source_id == other.source_id and
                    self.id == other.id and self.starred == other.starred)
        return NotImplemented

    def __init__(self, source, starred=True):
        self.source_id = source.id
        self.starred = starred

class InvalidUsernameException(Exception):

    """Raised when a user logs in with an invalid username"""

class LoginThrottledException(Exception):

    """Raised when a user attempts to log in
    too many times in a given time period"""

class WrongPasswordException(Exception):

    """Raised when a user logs in with an incorrect password"""

class BadTokenException(Exception):

    """Raised when a user logins in with an incorrect TOTP token"""

class PasswordError(Exception):

    """Generic error for passwords that are invalid.
    """

class InvalidPasswordLength(PasswordError):
    """Raised when attempting to create a Journalist or log in with an invalid
       password length.
    """

    def __init__(self, password):
        self.pw_len = len(password)

    def __str__(self):
        if self.pw_len > Journalist.MAX_PASSWORD_LEN:
            return "Password too long (len={})".format(self.pw_len)
        if self.pw_len < Journalist.MIN_PASSWORD_LEN:
            return "Password needs to be at least {} characters".format(
                Journalist.MIN_PASSWORD_LEN
            )

class NonDicewarePassword(PasswordError):

    """Raised when attempting to validate a password that is not diceware-like
    """

class Journalist(db.Model):
    __tablename__ = "journalists"
    id = Column(Integer, primary_key=True)
    username = Column(String(255), nullable=False, unique=True)
    pw_salt = Column(Binary(32))
    pw_hash = Column(Binary(256))
    is_admin = Column(Boolean)

    otp_secret = Column(String(16), default=pyotp.random_base32)
    is_totp = Column(Boolean, default=True)
    hotp_counter = Column(Integer, default=0)
    last_token = Column(String(6))

    created_on = Column(DateTime, default=datetime.datetime.utcnow)
    last_access = Column(DateTime)
    login_attempts = relationship(
        "JournalistLoginAttempt",
        backref="journalist")

    MIN_USERNAME_LEN = 3

    def __init__(self, username, password, is_admin=False, otp_secret=None):
        self.check_username_acceptable(username)
        self.username = username
        self.set_password(password)
        self.is_admin = is_admin

        if otp_secret:
            self.set_hotp_secret(otp_secret)

    def __repr__(self):
        return "<Journalist {0}{1}>".format(
            self.username,
            " [admin]" if self.is_admin else "")

    def _gen_salt(self, salt_bytes=32):
        return os.urandom(salt_bytes)

    _SCRYPT_PARAMS = dict(N=2**14, r=8, p=1)

    def _scrypt_hash(self, password, salt, params=None):
        if not params:
            params = self._SCRYPT_PARAMS
        return scrypt.hash(str(password), salt, **params)

    MAX_PASSWORD_LEN = 128
    MIN_PASSWORD_LEN = 14

    def set_password(self, password):
        self.check_password_acceptable(password)

        # Don't do anything if user's password hasn't changed.
        if self.pw_hash and self.valid_password(password):
            return

        self.pw_salt = self._gen_salt()
        self.pw_hash = self._scrypt_hash(password, self.pw_salt)

    @classmethod
    def check_username_acceptable(cls, username):
        if len(username) < cls.MIN_USERNAME_LEN:
            raise InvalidUsernameException(
                        'Username "{}" must be at least {} characters long.'
                        .format(username, cls.MIN_USERNAME_LEN))

    @classmethod
    def check_password_acceptable(cls, password):
        # Enforce a reasonable maximum length for passwords to avoid DoS
        if len(password) > cls.MAX_PASSWORD_LEN:
            raise InvalidPasswordLength(password)

        # Enforce a reasonable minimum length for new passwords
        if len(password) < cls.MIN_PASSWORD_LEN:
            raise InvalidPasswordLength(password)

        # Ensure all passwords are "diceware-like"
        if len(password.split()) < 7:
            raise NonDicewarePassword()

    def valid_password(self, password):
        # Avoid hashing passwords that are over the maximum length
        if len(password) > self.MAX_PASSWORD_LEN:
            raise InvalidPasswordLength(password)
        # No check on minimum password length here because some passwords
        # may have been set prior to setting the minimum password length.
        return pyotp.utils.compare_digest(
            self._scrypt_hash(password, self.pw_salt),
            self.pw_hash)

    def regenerate_totp_shared_secret(self):
        self.otp_secret = pyotp.random_base32()

    def set_hotp_secret(self, otp_secret):
        self.otp_secret = base64.b32encode(
            binascii.unhexlify(
                otp_secret.replace(
                    " ",
                    "")))
        self.is_totp = False
        self.hotp_counter = 0

    @property
    def totp(self):
        if self.is_totp:
            return pyotp.TOTP(self.otp_secret)
        else:
            raise ValueError('{} is not using TOTP'.format(self))

    @property
    def hotp(self):
        if not self.is_totp:
            return pyotp.HOTP(self.otp_secret)
        else:
            raise ValueError('{} is not using HOTP'.format(self))

    @property
    def shared_secret_qrcode(self):
        uri = self.totp.provisioning_uri(
            self.username,
            issuer_name="SecureDrop")

        qr = qrcode.QRCode(
            box_size=15,
            image_factory=qrcode.image.svg.SvgPathImage
        )
        qr.add_data(uri)
        img = qr.make_image()

        svg_out = StringIO()
        img.save(svg_out)
        return Markup(svg_out.getvalue())

    @property
    def formatted_otp_secret(self):
        """The OTP secret is easier to read and manually enter if it is all
        lowercase and split into four groups of four characters. The secret is
        base32-encoded, so it is case insensitive."""
        sec = self.otp_secret
        chunks = [sec[i:i + 4] for i in range(0, len(sec), 4)]
        return ' '.join(chunks).lower()

    def _format_token(self, token):
        """Strips from authentication tokens the whitespace
        that many clients add for readability"""
        return ''.join(token.split())

    def verify_token(self, token):
        token = self._format_token(token)

        # Store latest token to prevent OTP token reuse
        self.last_token = token
        db.session.commit()

        if self.is_totp:
            # Also check the given token against the previous and next
            # valid tokens, to compensate for potential time skew
            # between the client and the server. The total valid
            # window is 1:30s.
            return self.totp.verify(token, valid_window=1)
        else:
            for counter_val in range(
                    self.hotp_counter,
                    self.hotp_counter + 20):
                if self.hotp.verify(token, counter_val):
                    self.hotp_counter = counter_val + 1
                    db.session.commit()
                    return True
            return False

    _LOGIN_ATTEMPT_PERIOD = 60  # seconds
    _MAX_LOGIN_ATTEMPTS_PER_PERIOD = 5

    @classmethod
    def throttle_login(cls, user):
        # Record the login attempt...
        login_attempt = JournalistLoginAttempt(user)
        db.session.add(login_attempt)
        db.session.commit()

        # ...and reject it if they have exceeded the threshold
        login_attempt_period = datetime.datetime.utcnow() - \
            datetime.timedelta(seconds=cls._LOGIN_ATTEMPT_PERIOD)
        attempts_within_period = JournalistLoginAttempt.query.filter(
            JournalistLoginAttempt.journalist_id == user.id).filter(
            JournalistLoginAttempt.timestamp > login_attempt_period).all()
        if len(attempts_within_period) > cls._MAX_LOGIN_ATTEMPTS_PER_PERIOD:
            raise LoginThrottledException(
                "throttled ({} attempts in last {} seconds)".format(
                    len(attempts_within_period),
                    cls._LOGIN_ATTEMPT_PERIOD))

    @classmethod
    def login(cls, username, password, token):
        try:
            user = Journalist.query.filter_by(username=username).one()
        except NoResultFound:
            raise InvalidUsernameException(
                "invalid username '{}'".format(username))

        if LOGIN_HARDENING:
            cls.throttle_login(user)

            # Prevent TOTP token reuse
            if user.last_token is not None:
                if pyotp.utils.compare_digest(token, user.last_token):
                    raise BadTokenException("previously used token "
                                            "{}".format(token))
        if not user.verify_token(token):
            raise BadTokenException("invalid token")
        if not user.valid_password(password):
            raise WrongPasswordException("invalid password")
        return user

class JournalistLoginAttempt(db.Model):

    """This model keeps track of journalist's login attempts so we can
    rate limit them in order to prevent attackers from brute forcing
    passwords or two-factor tokens."""
    __tablename__ = "journalist_login_attempt"
    id = Column(Integer, primary_key=True)
    timestamp = Column(DateTime, default=datetime.datetime.utcnow)
    journalist_id = Column(Integer, ForeignKey('journalists.id'))

    def __init__(self, journalist):
        self.journalist_id = journalist.id
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
__author__ = "Jake Schurch"

import SecCrawler as sc
import re
import bs4
from bs4 import BeautifulSoup

def GetItemsFromSGML(SgmlHead):
    items = re.findall(r'(?:\<ITEMS\>)(\d{1}\.\d{2})', SgmlHead)
    return items

def CleanHtmlMarkup(markup: bs4.element.ResultSet):
    markup = markup.replace('', " ")
    markup = markup.replace("
", "\n
")
    markup = markup.replace("Table of Contents", "")

    return markup

def CleanTextMarkup(text: str):
    text = text.replace('\u200b', "")
    text = text.replace('• \n', '•')
    text = re.sub(r'[^\x00-\x7f]', r' ', text)
    text = re.sub(r'^\d+ $', r'', text, flags=re.MULTILINE)
    text = re.sub(r'\n\s+\n\s+\n\s+', '\n\n', text)

    return text

def GetCleanTags(HTML_Entity: bs4.element.ResultSet):
    FoundTags = HTML_Entity.find_all(True, recursive=False)
    for _tag in FoundTags:
        if _tag.name == 'table':
            for descendant in _tag.descendants:
                try:
                    if ('style' in descendant.attrs and
                            'background-color' in descendant.attrs['style']):
                        _tag.clear()
                        _tag.decompose()
                        break
                except AttributeError:
                    pass

        elif (_tag.attrs is not None and
              'style' in _tag.attrs and
              'text-align:center' in _tag.attrs['style']):

            if _tag.string == '' or _tag.string is None:
                pass
            else:
                _tag.clear()
                _tag.decompose()

        elif _tag.name == 'hr':
            _tag.clear()
            _tag.decompose()

        else:
            pass

    return FoundTags

def FilingTextParser(filing: sc.Filing):
    _Content = BeautifulSoup(filing.FilingText, "lxml").find(['document',
                                                             'text'])
    _CleanTags = GetCleanTags(_Content)

    _CleanHtml = CleanHtmlMarkup(
        ''.join(str(_tag) for _tag in _CleanTags))

    CleanText = CleanTextMarkup(
        BeautifulSoup(_CleanHtml, 'lxml').text)

    with open('parseFilingTesting.txt', 'w+') as f:
        f.write(CleanText)

    return CleanText

if __name__ == "__main__":
    testListing = [sc.Filing("goog", "10-k", totalFilingsWanted=1)]
    sc.SetInterimListing(testListing)

    FilingTextParser(sc.G_filingListing[0])
<commit_msg>Testing funciton annotaitons<commit_after># -*- coding: utf-8 -*-
__author__ = "Jake Schurch"

import SecCrawler as sc
import re
import bs4
from bs4 import BeautifulSoup
import copy
from datetime import datetime

def CheckInputTypes(func):
    def wrapper(*args, **kwargs):

        for arg in args:
            ArgValue = next(iter(locals()['args']))
            ArgParamType = next(iter(func.__annotations__.keys()))

            if type(ArgValue) != func.__annotations__[ArgParamType]:
                raise TypeError(
                    "Arg: '{0}' in {1}() is not of type: {2}".format(
                        ArgParamType, func.__name__,
                        func.__annotations__[ArgParamType]))

        for kwarg in kwargs:
            if type(locals()['kwargs'][kwarg]) != func.__annotations__[kwarg]:
                raise TypeError(
                    "Kwarg: '{0}' in {1}() is not of type: {2}").format(
                        kwarg, func.__name__,
                        func.__annotations__[kwarg])
        return func(*args, **kwargs)
    return wrapper

@CheckInputTypes
def foo(bar: int, buz: int):
    int3 = bar + buz
    return int3

foo("hello", "world")

def GetItemsFromSGML(SgmlHead):
    items = re.findall(r'(?:\<ITEMS\>)(\d{1}\.\d{2})', SgmlHead)
    return items

def CleanHtmlMarkup(markup: bs4.element.ResultSet):
    markup = markup.replace('', " ")
    markup = markup.replace("
", "\n
")
    markup = markup.replace("Table of Contents", "")

    return markup

def CleanTextMarkup(text: str):
    text = text.replace('\u200b', "")
    text = text.replace('• \n', '•')
    text = re.sub(r'[^\x00-\x7f]', r' ', text)
    text = re.sub(r'^\d+ $', r'', text, flags=re.MULTILINE)
    text = re.sub(r'\n\s+\n\s+\n\s+', '\n\n', text)

    return text

def GetCleanTags(HTML_Entity: bs4.element.ResultSet) -> 'ResultSet \
        not containing heavily formatted html tags':

    FoundTags = HTML_Entity.find_all(True, recursive=False)
    for _tag in FoundTags:
        if _tag.name == 'table':
            for descendant in _tag.descendants:
                try:
                    if ('style' in descendant.attrs and
                            'background-color' in descendant.attrs['style']):
                        _tag.clear()
                        _tag.decompose()
                        break
                except AttributeError:
                    pass

        elif (_tag.attrs is not None and
              'style' in _tag.attrs and
              'text-align:center' in _tag.attrs['style']):

            if _tag.string == '' or _tag.string is None:
                pass
            else:
                _tag.clear()
                _tag.decompose()

        elif _tag.name == 'hr':
            _tag.clear()
            _tag.decompose()

        else:
            pass

    return FoundTags

def FilingTextParser(filing: sc.Filing) -> 'raw filing text':

    _Content = BeautifulSoup(filing.FilingText, "lxml").find(['document',
                                                             'text'])
    _CleanTags = GetCleanTags(HTML_Entity=_Content)

    _CleanHtml = CleanHtmlMarkup(
        ''.join(str(_tag) for _tag in _CleanTags))

    CleanText = CleanTextMarkup(
        BeautifulSoup(_CleanHtml, 'lxml').text)

    with open('parseFilingTesting.txt', 'w+') as f:
        f.write(CleanText)

    return CleanText

def func():
    pass

'''d = globals().copy()

print(datetime.datetime.now())

for (item, val) in d.items():
    if type(d[item]) == type(func):
        globals()[item] = CheckInputTypes(d[item])

foo(bar='hello', buz='world')'''
<|endoftext|>