<commit_before>import os
import sys
import time
import json
import shlex
import base64
import sqlite3
import tempfile
import argparse
import webbrowser
import subprocess

from io import BytesIO
from pathlib import PurePosixPath, Path
from tempfile import NamedTemporaryFile
from xml.etree import ElementTree
from collections import namedtuple
from urllib.parse import quote, urlencode
from urllib.request import pathname2url

Account = namedtuple('Account', ['name', 'digits', 'period', 'secret', 'type', 'algorithm'])

def adb_list_dir(path):
    print('Listing directory', path, file=sys.stderr)

    # `adb exec-out` doesn't work properly on some devices. We have to fall back to `adb shell`,
    # which takes at least 600ms to exit even if the actual command runs quickly.
    # Reading a unique, non-existent file prints a predictable error message that delimits the end of
    # the stream, allowing us to let `adb shell` finish up its stuff in the background.
    lines = []
    process = subprocess.Popen(
        args=['adb', 'shell', f'su -c "ls -1 {shlex.quote(str(path))}; ls /3bb22bb739c29e435151cb38"'],
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT
    )

    for line in process.stdout:
        if b'ls: ' not in line:
            lines.append(line)
            continue

        message = line.partition(b'ls: ')[2].strip()
        process.kill()

        if b'3bb22bb739c29e435151cb38' in message:
            return [path/l[:-1].decode('utf-8') for l in lines]
        elif b'No such file or directory' in message:
            raise FileNotFoundError(path)
        else:
            raise IOError(message)

def adb_read_file(path):
    print('Reading file', path, file=sys.stderr)

    # `adb exec-out` doesn't work properly on some devices. We have to fall back to `adb shell`,
    # which takes at least 600ms to exit even if the actual command runs quickly.
    # Reading a unique, non-existent file prints a predictable error message that delimits the end of
    # the stream, allowing us to let `adb shell` finish up its stuff in the background.
    lines = []
    process = subprocess.Popen(
        args=['adb', 'shell', f'su -c "toybox base64 {shlex.quote(str(path))} /3bb22bb739c29e435151cb38"'],
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT
    )

    for line in process.stdout:
        if b'base64: ' not in line:
            lines.append(line)
            continue

        message = line.partition(b'base64: ')[2].strip()
        process.kill()

        if b'3bb22bb739c29e435151cb38' in message:
            return BytesIO(base64.b64decode(b''.join(lines)))
        elif b'No such file or directory' in message:
            raise FileNotFoundError(path)
        else:
            raise IOError(message)

def otpauth_encode_account(account):
    return f'otpauth://totp/{quote(account.name)}?' + urlencode({
        'secret': account.secret,
        'digits': account.digits,
        'period': account.period,
        'type': account.type,
        'algorithm': account.algorithm,
    })

def normalize_secret(secret):
    if set(secret.lower()) <= set('0123456789abcdef'):
        return base64.b32encode(bytes.fromhex(secret)).decode('ascii').rstrip('=')
    else:
        return secret.upper().rstrip('=')

def read_authy_accounts(data_root):
    for pref_file in ['com.authy.storage.tokens.authenticator.xml', 'com.authy.storage.tokens.authy.xml']:
        try:
            handle = adb_read_file(data_root/'com.authy.authy/shared_prefs'/pref_file)
        except FileNotFoundError:
            continue

        accounts = json.loads(ElementTree.parse(handle).find('string').text)

        for account in accounts:
            if 'decryptedSecret' in account:
                period = 30
                secret = account['decryptedSecret']
            else:
                period = 10
                secret = account['secretSeed']

            yield Account(account['name'], account['digits'], period, normalize_secret(secret), 'TOTP', 'SHA1')

def read_freeotp_accounts(data_root):
    try:
        handle = adb_read_file(data_root/'org.fedorahosted.freeotp/shared_prefs/tokens.xml')
    except FileNotFoundError:
        return

    for string in ElementTree.parse(handle).findall('string'):
        account = json.loads(string.text)

        if 'secret' not in account:
            continue

        if account['type'] != 'TOTP':
            raise ValueError('Only TOTP is supported.')

        secret = normalize_secret(bytes([b & 0xff for b in account['secret']]).hex())

        yield Account(account['label'], account['digits'], account['period'], secret, 'TOTP', 'SHA1')

def read_duo_accounts(data_root):
    try:
        handle = adb_read_file(data_root/'com.duosecurity.duomobile/files/duokit/accounts.json')
    except FileNotFoundError:
        return

    for account in json.load(handle):
        yield Account(account['name'], 6, 30, normalize_secret(account['otpGenerator']['otpSecret']), 'TOTP', 'SHA1')

def read_google_authenticator_accounts(data_root):
    try:
        database = adb_read_file(data_root/'com.google.android.apps.authenticator2/databases/databases')
    except FileNotFoundError:
        return

    with NamedTemporaryFile(delete=False, suffix='.html') as temp_handle:
        temp_handle.write(database.read())

    try:
        connection = sqlite3.connect(temp_handle.name)
        cursor = connection.cursor()
        cursor.execute('SELECT email, secret FROM accounts;')

        for name, secret in cursor.fetchall():
            yield Account(name, 6, 30, normalize_secret(secret), 'TOTP', 'SHA1')
    finally:
        os.unlink(temp_handle.name)

def read_microsoft_authenticator_accounts(data_root):
    try:
        database = adb_read_file(data_root/'com.azure.authenticator/databases/PhoneFactor')
    except FileNotFoundError:
        return

    with NamedTemporaryFile(delete=False) as temp_handle:
        temp_handle.write(database.read())

    try:
        connection = sqlite3.connect(temp_handle.name)
        cursor = connection.cursor()
        cursor.execute('SELECT name, oath_secret_key FROM accounts WHERE account_type=0;')

        for name, secret in cursor.fetchall():
            yield Account(name, 6, 30, normalize_secret(secret), 'TOTP', 'SHA1')
    finally:
        os.unlink(temp_handle.name)

def read_andotp_accounts(data_root):
    print('''
AndOTP encrypts its database with a key stored in the Android Keystore.
There is currently no way to interface with the Android Keystore and decrypt its database using just `adb` commands.
You will have to manually create a backup with AndOTP, which I can then read.
''')

    subprocess.check_output(['adb', 'shell', 'su', '-c', "am start -n org.shadowice.flocke.andotp/.Activities.BackupActivity"])
    backup_path = input('Enter the backup path (default: /sdcard/Download/otp_accounts.json): ') or '/sdcard/Download/otp_accounts.json'

    delete = input('Do you want to delete the backup afterwards (default: no)? ').lower() in ('yes', 'y')

    try:
        backup_data = adb_read_file(backup_path)
    except FileNotFoundError:
        print('Invalid path!', file=sys.stderr)
        return

    backup = json.load(backup_data)

    for account in backup:
        yield Account(account['label'], account['digits'], account['period'], normalize_secret(account['secret']), account['type'], account['algorithm'])

    if delete:
        subprocess.check_output(['adb', 'shell', 'rm', backup])

def read_steam_authenticator_accounts(data_root):
    accounts_folder = 'com.valvesoftware.android.steam.community/files'
    try:
        account_files = adb_list_dir(data_root/accounts_folder)
    except FileNotFoundError:
        return

    for account_file in account_files:
        account_json = json.load(adb_read_file(account_file))

        secret = normalize_secret(base64.b32encode(base64.b64decode(account_json['shared_secret'])))

        yield Account(account_json['account_name'], 5, 30, secret, 'STEAM', 'SHA1')

def export_andotp(accounts):
    return json.dumps([{
        'secret': a.secret if not a.name.startswith('steam-') else a.secret.decode('utf-8'),
        'label': a.name,
        'digits': a.digits,
        'period': a.period,
        'type': a.type,
        'algorithm': 'SHA1'
    } for a in accounts])

def display_qr_codes(accounts):
    accounts_html = '''
        <!doctype html>

        <meta http-equiv="Cache-Control" content="no-cache, no-store, must-revalidate" />
        <meta http-equiv="Pragma" content="no-cache" />
        <meta http-equiv="Expires" content="0" />

        <title>OTP QR Codes</title>
        <style type="text/css">
            body {
                width: 100%%;
            }
        </style>

        <body>
            <script src="https://cdnjs.cloudflare.com/ajax/libs/qrious/4.0.2/qrious.min.js" integrity="sha384-Dr98ddmUw2QkdCarNQ+OL7xLty7cSxgR0T7v1tq4UErS/qLV0132sBYTolRAFuOV" crossorigin="anonymous"></script>
            <script>
                var accounts = %s;

                for (var i = 0; i < accounts.length; i++) {
                    var account = accounts[i];

                    var heading = document.createElement('h2');
                    heading.textContent = decodeURIComponent(account.split('?')[0].split('/')[3]);
                    document.body.appendChild(heading);

                    var image = document.createElement('img');
                    image.style.width = '100%%';
                    image.style.maxWidth = '500px';
                    image.style.height = 'auto';
                    document.body.appendChild(image);

                    var qr_image = new QRious({
                        element: image,
                        value: account,
                        size: 500
                    });
                }
            </script>
        </body>''' % json.dumps([otpauth_encode_account(a) for a in accounts])

    # Temporary files are only readable by the current user (mode 0600)
    with tempfile.NamedTemporaryFile(delete=False) as handle:
        handle.write(accounts_html.encode('utf-8'))

    try:
        webbrowser.open(f'file:{pathname2url(handle.name)}')
        time.sleep(10)  # webbrowser.open exits immediately so we should wait before deleting the file
    finally:
        os.remove(handle.name)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Extracts TOTP secrets from a rooted Android phone.', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--andotp', action='store_true', help='parse AndOTP codes from a backup')
    parser.add_argument('--no-authy', action='store_true', help='no Authy codes')
    parser.add_argument('--no-duo', action='store_true', help='no Duo codes')
    parser.add_argument('--no-freeotp', action='store_true', help='no FreeOTP codes')
    parser.add_argument('--no-google-authenticator', action='store_true', help='no Google Authenticator codes')
    parser.add_argument('--no-microsoft-authenticator', action='store_true', help='no Microsoft Authenticator codes')
    parser.add_argument('--no-steam-authenticator', action='store_true', help='no Steam Authenticator codes')
    parser.add_argument('--data', type=PurePosixPath, default=PurePosixPath('/data/data/'), help='path to the app data folder')
    parser.add_argument('--no-show-uri', action='store_true', help='disable printing the accounts as otpauth:// URIs')
    parser.add_argument('--show-qr', action='store_true', help='displays the accounts as a local webpage with scannable QR codes')
    parser.add_argument('--andotp-backup', type=Path, help='saves the accounts as an AndOTP backup file')

    args = parser.parse_args()

    print(f'Checking for root by listing the contents of {args.data}. You might have to grant ADB temporary root access.', file=sys.stderr)

    if not adb_list_dir(args.data):
        print('Root not found or data directory is incorrect!', file=sys.stderr)
        sys.exit(1)

    accounts = set()

    if args.andotp:
        accounts.update(read_andotp_accounts(args.data))

    if not args.no_authy:
        accounts.update(read_authy_accounts(args.data))

    if not args.no_duo:
        accounts.update(read_duo_accounts(args.data))

    if not args.no_freeotp:
        accounts.update(read_freeotp_accounts(args.data))

    if not args.no_google_authenticator:
        accounts.update(read_google_authenticator_accounts(args.data))

    if not args.no_microsoft_authenticator:
        accounts.update(read_microsoft_authenticator_accounts(args.data))

    if not args.no_steam_authenticator:
        accounts.update(read_steam_authenticator_accounts(args.data))

    if not args.no_show_uri:
        for account in accounts:
            print(otpauth_encode_account(account))

    if args.show_qr:
        display_qr_codes(accounts)

    if args.andotp_backup:
        with open(args.andotp_backup, 'wb') as handle:
            handle.write(export_andotp(accounts).encode('utf-8'))
<commit_msg>Closed the database connection after we're through with it. Fixes #10<commit_after>import os
import sys
import time
import json
import shlex
import base64
import sqlite3
import tempfile
import argparse
import webbrowser
import subprocess

from io import BytesIO
from pathlib import PurePosixPath, Path
from tempfile import NamedTemporaryFile
from xml.etree import ElementTree
from collections import namedtuple
from urllib.parse import quote, urlencode
from urllib.request import pathname2url

Account = namedtuple('Account', ['name', 'digits', 'period', 'secret', 'type', 'algorithm'])

def adb_list_dir(path):
    print('Listing directory', path, file=sys.stderr)

    # `adb exec-out` doesn't work properly on some devices. We have to fall back to `adb shell`,
    # which takes at least 600ms to exit even if the actual command runs quickly.
    # Reading a unique, non-existent file prints a predictable error message that delimits the end of
    # the stream, allowing us to let `adb shell` finish up its stuff in the background.
    lines = []
    process = subprocess.Popen(
        args=['adb', 'shell', f'su -c "ls -1 {shlex.quote(str(path))}; ls /3bb22bb739c29e435151cb38"'],
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT
    )

    for line in process.stdout:
        if b'ls: ' not in line:
            lines.append(line)
            continue

        message = line.partition(b'ls: ')[2].strip()
        process.kill()

        if b'3bb22bb739c29e435151cb38' in message:
            return [path/l[:-1].decode('utf-8') for l in lines]
        elif b'No such file or directory' in message:
            raise FileNotFoundError(path)
        else:
            raise IOError(message)

def adb_read_file(path):
    print('Reading file', path, file=sys.stderr)

    # `adb exec-out` doesn't work properly on some devices. We have to fall back to `adb shell`,
    # which takes at least 600ms to exit even if the actual command runs quickly.
    # Reading a unique, non-existent file prints a predictable error message that delimits the end of
    # the stream, allowing us to let `adb shell` finish up its stuff in the background.
    lines = []
    process = subprocess.Popen(
        args=['adb', 'shell', f'su -c "toybox base64 {shlex.quote(str(path))} /3bb22bb739c29e435151cb38"'],
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT
    )

    for line in process.stdout:
        if b'base64: ' not in line:
            lines.append(line)
            continue

        message = line.partition(b'base64: ')[2].strip()
        process.kill()

        if b'3bb22bb739c29e435151cb38' in message:
            return BytesIO(base64.b64decode(b''.join(lines)))
        elif b'No such file or directory' in message:
            raise FileNotFoundError(path)
        else:
            raise IOError(message)

def otpauth_encode_account(account):
    return f'otpauth://totp/{quote(account.name)}?' + urlencode({
        'secret': account.secret,
        'digits': account.digits,
        'period': account.period,
        'type': account.type,
        'algorithm': account.algorithm,
    })

def normalize_secret(secret):
    if set(secret.lower()) <= set('0123456789abcdef'):
        return base64.b32encode(bytes.fromhex(secret)).decode('ascii').rstrip('=')
    else:
        return secret.upper().rstrip('=')

def read_authy_accounts(data_root):
    for pref_file in ['com.authy.storage.tokens.authenticator.xml', 'com.authy.storage.tokens.authy.xml']:
        try:
            handle = adb_read_file(data_root/'com.authy.authy/shared_prefs'/pref_file)
        except FileNotFoundError:
            continue

        accounts = json.loads(ElementTree.parse(handle).find('string').text)

        for account in accounts:
            if 'decryptedSecret' in account:
                period = 30
                secret = account['decryptedSecret']
            else:
                period = 10
                secret = account['secretSeed']

            yield Account(account['name'], account['digits'], period, normalize_secret(secret), 'TOTP', 'SHA1')

def read_freeotp_accounts(data_root):
    try:
        handle = adb_read_file(data_root/'org.fedorahosted.freeotp/shared_prefs/tokens.xml')
    except FileNotFoundError:
        return

    for string in ElementTree.parse(handle).findall('string'):
        account = json.loads(string.text)

        if 'secret' not in account:
            continue

        if account['type'] != 'TOTP':
            raise ValueError('Only TOTP is supported.')

        secret = normalize_secret(bytes([b & 0xff for b in account['secret']]).hex())

        yield Account(account['label'], account['digits'], account['period'], secret, 'TOTP', 'SHA1')

def read_duo_accounts(data_root):
    try:
        handle = adb_read_file(data_root/'com.duosecurity.duomobile/files/duokit/accounts.json')
    except FileNotFoundError:
        return

    for account in json.load(handle):
        yield Account(account['name'], 6, 30, normalize_secret(account['otpGenerator']['otpSecret']), 'TOTP', 'SHA1')

def read_google_authenticator_accounts(data_root):
    try:
        database = adb_read_file(data_root/'com.google.android.apps.authenticator2/databases/databases')
    except FileNotFoundError:
        return

    with NamedTemporaryFile(delete=False, suffix='.html') as temp_handle:
        temp_handle.write(database.read())

    try:
        connection = sqlite3.connect(temp_handle.name)
        cursor = connection.cursor()
        cursor.execute('SELECT email, secret FROM accounts;')

        for name, secret in cursor.fetchall():
            yield Account(name, 6, 30, normalize_secret(secret), 'TOTP', 'SHA1')

        connection.close()
    finally:
        os.unlink(temp_handle.name)

def read_microsoft_authenticator_accounts(data_root):
    try:
        database = adb_read_file(data_root/'com.azure.authenticator/databases/PhoneFactor')
    except FileNotFoundError:
        return

    with NamedTemporaryFile(delete=False) as temp_handle:
        temp_handle.write(database.read())

    try:
        connection = sqlite3.connect(temp_handle.name)
        cursor = connection.cursor()
        cursor.execute('SELECT name, oath_secret_key FROM accounts WHERE account_type=0;')

        for name, secret in cursor.fetchall():
            yield Account(name, 6, 30, normalize_secret(secret), 'TOTP', 'SHA1')
    finally:
        os.unlink(temp_handle.name)

def read_andotp_accounts(data_root):
    print('''
AndOTP encrypts its database with a key stored in the Android Keystore.
There is currently no way to interface with the Android Keystore and decrypt its database using just `adb` commands.
You will have to manually create a backup with AndOTP, which I can then read.
''')

    subprocess.check_output(['adb', 'shell', 'su', '-c', "am start -n org.shadowice.flocke.andotp/.Activities.BackupActivity"])
    backup_path = input('Enter the backup path (default: /sdcard/Download/otp_accounts.json): ') or '/sdcard/Download/otp_accounts.json'

    delete = input('Do you want to delete the backup afterwards (default: no)? ').lower() in ('yes', 'y')

    try:
        backup_data = adb_read_file(backup_path)
    except FileNotFoundError:
        print('Invalid path!', file=sys.stderr)
        return

    backup = json.load(backup_data)

    for account in backup:
        yield Account(account['label'], account['digits'], account['period'], normalize_secret(account['secret']), account['type'], account['algorithm'])

    if delete:
        subprocess.check_output(['adb', 'shell', 'rm', backup])

def read_steam_authenticator_accounts(data_root):
    accounts_folder = 'com.valvesoftware.android.steam.community/files'
    try:
        account_files = adb_list_dir(data_root/accounts_folder)
    except FileNotFoundError:
        return

    for account_file in account_files:
        account_json = json.load(adb_read_file(account_file))

        secret = normalize_secret(base64.b32encode(base64.b64decode(account_json['shared_secret'])))

        yield Account(account_json['account_name'], 5, 30, secret, 'STEAM', 'SHA1')

def export_andotp(accounts):
    return json.dumps([{
        'secret': a.secret if not a.name.startswith('steam-') else a.secret.decode('utf-8'),
        'label': a.name,
        'digits': a.digits,
        'period': a.period,
        'type': a.type,
        'algorithm': 'SHA1'
    } for a in accounts])

def display_qr_codes(accounts):
    accounts_html = '''
        <!doctype html>

        <meta http-equiv="Cache-Control" content="no-cache, no-store, must-revalidate" />
        <meta http-equiv="Pragma" content="no-cache" />
        <meta http-equiv="Expires" content="0" />

        <title>OTP QR Codes</title>
        <style type="text/css">
            body {
                width: 100%%;
            }
        </style>

        <body>
            <script src="https://cdnjs.cloudflare.com/ajax/libs/qrious/4.0.2/qrious.min.js" integrity="sha384-Dr98ddmUw2QkdCarNQ+OL7xLty7cSxgR0T7v1tq4UErS/qLV0132sBYTolRAFuOV" crossorigin="anonymous"></script>
            <script>
                var accounts = %s;

                for (var i = 0; i < accounts.length; i++) {
                    var account = accounts[i];

                    var heading = document.createElement('h2');
                    heading.textContent = decodeURIComponent(account.split('?')[0].split('/')[3]);
                    document.body.appendChild(heading);

                    var image = document.createElement('img');
                    image.style.width = '100%%';
                    image.style.maxWidth = '500px';
                    image.style.height = 'auto';
                    document.body.appendChild(image);

                    var qr_image = new QRious({
                        element: image,
                        value: account,
                        size: 500
                    });
                }
            </script>
        </body>''' % json.dumps([otpauth_encode_account(a) for a in accounts])

    # Temporary files are only readable by the current user (mode 0600)
    with tempfile.NamedTemporaryFile(delete=False) as handle:
        handle.write(accounts_html.encode('utf-8'))

    try:
        webbrowser.open(f'file:{pathname2url(handle.name)}')
        time.sleep(10)  # webbrowser.open exits immediately so we should wait before deleting the file
    finally:
        os.remove(handle.name)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Extracts TOTP secrets from a rooted Android phone.', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--andotp', action='store_true', help='parse AndOTP codes from a backup')
    parser.add_argument('--no-authy', action='store_true', help='no Authy codes')
    parser.add_argument('--no-duo', action='store_true', help='no Duo codes')
    parser.add_argument('--no-freeotp', action='store_true', help='no FreeOTP codes')
    parser.add_argument('--no-google-authenticator', action='store_true', help='no Google Authenticator codes')
    parser.add_argument('--no-microsoft-authenticator', action='store_true', help='no Microsoft Authenticator codes')
    parser.add_argument('--no-steam-authenticator', action='store_true', help='no Steam Authenticator codes')
    parser.add_argument('--data', type=PurePosixPath, default=PurePosixPath('/data/data/'), help='path to the app data folder')
    parser.add_argument('--no-show-uri', action='store_true', help='disable printing the accounts as otpauth:// URIs')
    parser.add_argument('--show-qr', action='store_true', help='displays the accounts as a local webpage with scannable QR codes')
    parser.add_argument('--andotp-backup', type=Path, help='saves the accounts as an AndOTP backup file')

    args = parser.parse_args()

    print(f'Checking for root by listing the contents of {args.data}. You might have to grant ADB temporary root access.', file=sys.stderr)

    if not adb_list_dir(args.data):
        print('Root not found or data directory is incorrect!', file=sys.stderr)
        sys.exit(1)

    accounts = set()

    if args.andotp:
        accounts.update(read_andotp_accounts(args.data))

    if not args.no_authy:
        accounts.update(read_authy_accounts(args.data))

    if not args.no_duo:
        accounts.update(read_duo_accounts(args.data))

    if not args.no_freeotp:
        accounts.update(read_freeotp_accounts(args.data))

    if not args.no_google_authenticator:
        accounts.update(read_google_authenticator_accounts(args.data))

    if not args.no_microsoft_authenticator:
        accounts.update(read_microsoft_authenticator_accounts(args.data))

    if not args.no_steam_authenticator:
        accounts.update(read_steam_authenticator_accounts(args.data))

    if not args.no_show_uri:
        for account in accounts:
            print(otpauth_encode_account(account))

    if args.show_qr:
        display_qr_codes(accounts)

    if args.andotp_backup:
        with open(args.andotp_backup, 'wb') as handle:
            handle.write(export_andotp(accounts).encode('utf-8'))
<|endoftext|>


<commit_before>import pytest

from pytest_bdd import scenario, given, when, then
from pytest_bdd.scenario import GivenAlreadyUsed

test_steps = scenario('steps.feature', 'Executed step by step')

@given('I have a foo fixture with value "foo"')
def foo():
    return 'foo'

@given('there is a list')
def results():
    return []

@when('I append 1 to the list')
def append_1(results):
    results.append(1)

@when('I append 2 to the list')
def append_2(results):
    results.append(2)

@when('I append 3 to the list')
def append_3(results):
    results.append(3)

@then('foo should have value "foo"')
def foo_is_foo(foo):
    assert foo == 'foo'

@then('the list should be [1, 2, 3]')
def check_results(results):
    assert results == [1, 2, 3]

test_when_first = scenario('steps.feature', 'When step can be the first')

@when('I do nothing')
def do_nothing():
    pass

@then('I make no mistakes')
def no_errors():
    assert True

test_then_after_given = scenario('steps.feature', 'Then step can follow Given step')

@given('xyz')
def xyz():
    """Used in the test_same_step_name."""
    return

test_conftest = scenario('steps.feature', 'All steps are declared in the conftest')

def test_multiple_given(request):
    """Using the same given fixture raises an error."""
    test = scenario(
        'steps.feature',
        'Using the same given fixture raises an error',
    )
    with pytest.raises(GivenAlreadyUsed):
        test(request)

def test_step_hooks(testdir):
    """When step fails."""
    testdir.makefile('.feature', test="""
    Scenario: When step has hook on failure
        Given I have a bar
        When it fails

    Scenario: When step is not found
        Given not found

    Scenario: When step validation error happens
        Given foo
        And foo
    """)
    testdir.makepyfile("""
        from pytest_bdd import given, when, scenario

        def pytest_bdd_step_error():
            print 'ok'

        @given('I have a bar')
        def i_have_bar():
            return 'bar'

        @when('it fails')
        def when_it_fails():
            raise Exception('when fails')

        @scenario('test.feature', 'When step has hook on failure')
        def test_when_fails():
            pass

        @scenario('test.feature', 'When step is not found')
        def test_when_not_found():
            pass

        @when('foo')
        def foo():
            return 'foo'

        @scenario('test.feature', 'When step validation error happens')
        def test_when_step_validation_error():
            pass
    """)
    reprec = testdir.inline_run("-k test_when_fails")
    assert reprec.ret == 1

    calls = reprec.getcalls("pytest_bdd_step_func_lookup")
    assert calls[0].request

    calls = reprec.getcalls("pytest_bdd_step_func_found")
    assert calls[0].request

    calls = reprec.getcalls("pytest_bdd_step_start")
    assert calls[0].request

    calls = reprec.getcalls("pytest_bdd_step_finish")
    assert calls[0].request

    calls = reprec.getcalls("pytest_bdd_step_error")
    assert calls[0].request

    reprec = testdir.inline_run("-k test_when_not_found")
    assert reprec.ret == 3

    calls = reprec.getcalls("pytest_bdd_step_func_lookup_error")
    assert calls[0].request

    reprec = testdir.inline_run("-k test_when_step_validation_error")
    assert reprec.ret == 3
<commit_msg>remove print<commit_after>import pytest

from pytest_bdd import scenario, given, when, then
from pytest_bdd.scenario import GivenAlreadyUsed

test_steps = scenario('steps.feature', 'Executed step by step')

@given('I have a foo fixture with value "foo"')
def foo():
    return 'foo'

@given('there is a list')
def results():
    return []

@when('I append 1 to the list')
def append_1(results):
    results.append(1)

@when('I append 2 to the list')
def append_2(results):
    results.append(2)

@when('I append 3 to the list')
def append_3(results):
    results.append(3)

@then('foo should have value "foo"')
def foo_is_foo(foo):
    assert foo == 'foo'

@then('the list should be [1, 2, 3]')
def check_results(results):
    assert results == [1, 2, 3]

test_when_first = scenario('steps.feature', 'When step can be the first')

@when('I do nothing')
def do_nothing():
    pass

@then('I make no mistakes')
def no_errors():
    assert True

test_then_after_given = scenario('steps.feature', 'Then step can follow Given step')

@given('xyz')
def xyz():
    """Used in the test_same_step_name."""
    return

test_conftest = scenario('steps.feature', 'All steps are declared in the conftest')

def test_multiple_given(request):
    """Using the same given fixture raises an error."""
    test = scenario(
        'steps.feature',
        'Using the same given fixture raises an error',
    )
    with pytest.raises(GivenAlreadyUsed):
        test(request)

def test_step_hooks(testdir):
    """When step fails."""
    testdir.makefile('.feature', test="""
    Scenario: When step has hook on failure
        Given I have a bar
        When it fails

    Scenario: When step is not found
        Given not found

    Scenario: When step validation error happens
        Given foo
        And foo
    """)
    testdir.makepyfile("""
        from pytest_bdd import given, when, scenario

        @given('I have a bar')
        def i_have_bar():
            return 'bar'

        @when('it fails')
        def when_it_fails():
            raise Exception('when fails')

        @scenario('test.feature', 'When step has hook on failure')
        def test_when_fails():
            pass

        @scenario('test.feature', 'When step is not found')
        def test_when_not_found():
            pass

        @when('foo')
        def foo():
            return 'foo'

        @scenario('test.feature', 'When step validation error happens')
        def test_when_step_validation_error():
            pass
    """)
    reprec = testdir.inline_run("-k test_when_fails")
    assert reprec.ret == 1

    calls = reprec.getcalls("pytest_bdd_step_func_lookup")
    assert calls[0].request

    calls = reprec.getcalls("pytest_bdd_step_func_found")
    assert calls[0].request

    calls = reprec.getcalls("pytest_bdd_step_start")
    assert calls[0].request

    calls = reprec.getcalls("pytest_bdd_step_finish")
    assert calls[0].request

    calls = reprec.getcalls("pytest_bdd_step_error")
    assert calls[0].request

    reprec = testdir.inline_run("-k test_when_not_found")
    assert reprec.ret == 3

    calls = reprec.getcalls("pytest_bdd_step_func_lookup_error")
    assert calls[0].request

    reprec = testdir.inline_run("-k test_when_step_validation_error")
    assert reprec.ret == 3
<|endoftext|>


<commit_before>
# Local modules
from common import debug
from common import chrono
from common.telegram_utils.telegram_class import TelegramPost

from common.user.bibleuser_utils import *

from common.constants import *

TELEGRAM_MAX_LENGTH = 4096

TELEGRAM_OPTION_REPLY_KEYBOARD = 'reply_markup'
TELEGRAM_KEYBOARD_GRID_SIZE = 3

# Telegram message sending functionality
def format_keyboard(options=[]):
    num_buttons = len(options)
    modulus = 1 if num_buttons % TELEGRAM_KEYBOARD_GRID_SIZE else 0
    num_rows = int(num_buttons / TELEGRAM_KEYBOARD_GRID_SIZE) + modulus

    keyboard_data = []
    for i in range(0, num_rows):
        keyboard_row = []

        for j in range(0, TELEGRAM_KEYBOARD_GRID_SIZE):
            if num_buttons == 0:
                break

            data = options[i * TELEGRAM_KEYBOARD_GRID_SIZE + j]
            keyboard_row.append({'text': data})
            num_buttons -= 1

        keyboard_data.append(keyboard_row)

    return keyboard_data

def send_msg(msg, id):
    debug.log('Sending message to ' + str(id) + ': ' +  msg)

    last = None
    chunks = []
    while len(msg) > TELEGRAM_MAX_LENGTH:
        last = msg.rfind(' ', 0, TELEGRAM_MAX_LENGTH)
        if last == -1:
            last = TELEGRAM_MAX_LENGTH

        debug.log('Chunk: ' + msg[:last])
        chunks.append(msg[:last])
        msg = msg[last:]
        last = None

    chunks.append(msg[last:])

    for chunk in chunks:
        post = TelegramPost(id)
        post.add_text(chunk)
        post.send()

def send_msg_keyboard(msg, id, options=[], inline=False, one_time=False):
    post = TelegramPost(id)
    post.add_text(msg)
    if inline:
        post.add_inline_keyboard(options)
    else:
        post.add_keyboard(options, one_time)
    post.send()

def send_close_keyboard(msg, id):
    post = TelegramPost(id)
    post.add_text(msg)
    post.close_keyboard()
    post.send()

# Telegram message parsing
def parse_payload(msg):
    if msg is None:
        return None

    text = msg.get('text')
    if text is not None:
        return text

    audio = msg.get('audio')
    if audio is not None:
        return audio

    document = msg.get('document')
    if document is not None:
        return document

    photo = msg.get('photo')
    if photo is not None:
        return photo

    sticker = msg.get('sticker')
    if sticker is not None:
        return sticker

    video = msg.get('video')
    if video is not None:
        return video

    voice = msg.get('voice')
    if voice is not None:
        return voice

    return None

# Telegram message prettifying
def surround(text, front, back = None):
    if back is None:
        back = front

    return front + text + back

def bold(text):
    return surround(text, '*')

def italics(text):
    return surround(text, '_')

def bracket(text):
    return surround(text, '(', ')')

def bracket_square(text):
    return surround(text, '[', ']')

def link(text, link):
    return bracket_square(text) + bracket(link)

def join(blocks, separator):
    return separator.join(blocks)<commit_msg>Fixing typo in import<commit_after>
# Local modules
from common import debug
from common import chrono
from common.telegram.telegram_class import TelegramPost

from common.user.bibleuser_utils import *

from common.constants import *

TELEGRAM_MAX_LENGTH = 4096

TELEGRAM_OPTION_REPLY_KEYBOARD = 'reply_markup'
TELEGRAM_KEYBOARD_GRID_SIZE = 3

# Telegram message sending functionality
def format_keyboard(options=[]):
    num_buttons = len(options)
    modulus = 1 if num_buttons % TELEGRAM_KEYBOARD_GRID_SIZE else 0
    num_rows = int(num_buttons / TELEGRAM_KEYBOARD_GRID_SIZE) + modulus

    keyboard_data = []
    for i in range(0, num_rows):
        keyboard_row = []

        for j in range(0, TELEGRAM_KEYBOARD_GRID_SIZE):
            if num_buttons == 0:
                break

            data = options[i * TELEGRAM_KEYBOARD_GRID_SIZE + j]
            keyboard_row.append({'text': data})
            num_buttons -= 1

        keyboard_data.append(keyboard_row)

    return keyboard_data

def send_msg(msg, id):
    debug.log('Sending message to ' + str(id) + ': ' +  msg)

    last = None
    chunks = []
    while len(msg) > TELEGRAM_MAX_LENGTH:
        last = msg.rfind(' ', 0, TELEGRAM_MAX_LENGTH)
        if last == -1:
            last = TELEGRAM_MAX_LENGTH

        debug.log('Chunk: ' + msg[:last])
        chunks.append(msg[:last])
        msg = msg[last:]
        last = None

    chunks.append(msg[last:])

    for chunk in chunks:
        post = TelegramPost(id)
        post.add_text(chunk)
        post.send()

def send_msg_keyboard(msg, id, options=[], inline=False, one_time=False):
    post = TelegramPost(id)
    post.add_text(msg)
    if inline:
        post.add_inline_keyboard(options)
    else:
        post.add_keyboard(options, one_time)
    post.send()

def send_close_keyboard(msg, id):
    post = TelegramPost(id)
    post.add_text(msg)
    post.close_keyboard()
    post.send()

# Telegram message parsing
def parse_payload(msg):
    if msg is None:
        return None

    text = msg.get('text')
    if text is not None:
        return text

    audio = msg.get('audio')
    if audio is not None:
        return audio

    document = msg.get('document')
    if document is not None:
        return document

    photo = msg.get('photo')
    if photo is not None:
        return photo

    sticker = msg.get('sticker')
    if sticker is not None:
        return sticker

    video = msg.get('video')
    if video is not None:
        return video

    voice = msg.get('voice')
    if voice is not None:
        return voice

    return None

# Telegram message prettifying
def surround(text, front, back = None):
    if back is None:
        back = front

    return front + text + back

def bold(text):
    return surround(text, '*')

def italics(text):
    return surround(text, '_')

def bracket(text):
    return surround(text, '(', ')')

def bracket_square(text):
    return surround(text, '[', ']')

def link(text, link):
    return bracket_square(text) + bracket(link)

def join(blocks, separator):
    return separator.join(blocks)<|endoftext|>


<commit_before>#
# See top-level LICENSE.rst file for Copyright information
#
# -*- coding: utf-8 -*-

from __future__ import absolute_import, division, print_function

from collections import OrderedDict

from ..defs import (task_name_sep, task_state_to_int, task_int_to_state)

from ...util import option_list

from ...io import findfile

from .base import BaseTask

import sys,re,os,copy

# NOTE: only one class in this file should have a name that starts with "Task".

class TaskFluxCalib(BaseTask):
    """Class containing the properties of a sky fit task.
    """
    def __init__(self):
        super(TaskFluxCalib, self).__init__()
        # then put int the specifics of this class
        # _cols must have a state
        self._type = "fluxcalib"
        self._cols = [
            "night",
            "band",
            "spec",
            "expid",
            "state"
        ]
        self._coltypes = [
            "integer",
            "text",
            "integer",
            "integer",
            "integer"
        ]
        # _name_fields must also be in _cols
        self._name_fields  = ["night","band","spec","expid"]
        self._name_formats = ["08d","s","d","08d"]

    def _paths(self, name):
        """See BaseTask.paths.
        """
        props = self.name_split(name)
        camera = "{}{}".format(props["band"], props["spec"])
        return [ findfile("calib", night=props["night"], expid=props["expid"],
            camera=camera, groupname=None, nside=None, band=props["band"],
            spectrograph=props["spec"]) ]

    def _deps(self, name, db, inputs):
        """See BaseTask.deps.
        """
        from .base import task_classes
        props = self.name_split(name)
        deptasks = {
            "infile" : task_classes["extract"].name_join(props),
            "fiberflat" : task_classes["fiberflatnight"].name_join(props),
            "sky" : task_classes["sky"].name_join(props),
            "models" : task_classes["starfit"].name_join(props)
        }
        return deptasks

    def _run_max_procs(self, procs_per_node):
        """See BaseTask.run_max_procs.
        """
        return 1

    def _run_time(self, name, procs_per_node, db=None):
        """See BaseTask.run_time.
        """
        return 2 # less than a minute (for the simple sky fit)

    def _run_defaults(self):
        """See BaseTask.run_defaults.
        """
        opts = {}
        return opts

    def _option_list(self, name, opts):
        """Build the full list of options.

        This includes appending the filenames and incorporating runtime
        options.
        """
        from .base import task_classes, task_type

        deps = self.deps(name)
        options = {}
        options["infile"]    = task_classes["extract"].paths(deps["infile"])[0]
        options["fiberflat"] = task_classes["fiberflatnight"].paths(deps["fiberflat"])[0]
        options["sky"]       = task_classes["sky"].paths(deps["sky"])[0]
        options["models"]    = task_classes["starfit"].paths(deps["models"])[0]
        options["outfile"]   = self.paths(name)[0]

        options.update(opts)
        return option_list(options)

    def _run_cli(self, name, opts, procs, db=None):
        """See BaseTask.run_cli.
        """
        entry = "desi_compute_fluxcalibration"
        optlist = self._option_list(name, opts)
        com = "{} {}".format(entry, " ".join(optlist))
        return com

    def _run(self, name, opts, comm, db=None):
        """See BaseTask.run.
        """
        from ...scripts import fluxcalibration
        optlist = self._option_list(name, opts)
        args = fluxcalibration.parse(optlist)
        fluxcalibration.main(args)
        return
<commit_msg>automatic getready<commit_after>#
# See top-level LICENSE.rst file for Copyright information
#
# -*- coding: utf-8 -*-

from __future__ import absolute_import, division, print_function

from collections import OrderedDict

from ..defs import (task_name_sep, task_state_to_int, task_int_to_state)

from ...util import option_list

from ...io import findfile

from .base import (BaseTask, task_classes)

from desiutil.log import get_logger

import sys,re,os,copy

# NOTE: only one class in this file should have a name that starts with "Task".

class TaskFluxCalib(BaseTask):
    """Class containing the properties of a sky fit task.
    """
    def __init__(self):
        super(TaskFluxCalib, self).__init__()
        # then put int the specifics of this class
        # _cols must have a state
        self._type = "fluxcalib"
        self._cols = [
            "night",
            "band",
            "spec",
            "expid",
            "state"
        ]
        self._coltypes = [
            "integer",
            "text",
            "integer",
            "integer",
            "integer"
        ]
        # _name_fields must also be in _cols
        self._name_fields  = ["night","band","spec","expid"]
        self._name_formats = ["08d","s","d","08d"]

    def _paths(self, name):
        """See BaseTask.paths.
        """
        props = self.name_split(name)
        camera = "{}{}".format(props["band"], props["spec"])
        return [ findfile("calib", night=props["night"], expid=props["expid"],
            camera=camera, groupname=None, nside=None, band=props["band"],
            spectrograph=props["spec"]) ]

    def _deps(self, name, db, inputs):
        """See BaseTask.deps.
        """
        from .base import task_classes
        props = self.name_split(name)
        deptasks = {
            "infile" : task_classes["extract"].name_join(props),
            "fiberflat" : task_classes["fiberflatnight"].name_join(props),
            "sky" : task_classes["sky"].name_join(props),
            "models" : task_classes["starfit"].name_join(props)
        }
        return deptasks

    def _run_max_procs(self, procs_per_node):
        """See BaseTask.run_max_procs.
        """
        return 1

    def _run_time(self, name, procs_per_node, db=None):
        """See BaseTask.run_time.
        """
        return 2 # less than a minute (for the simple sky fit)

    def _run_defaults(self):
        """See BaseTask.run_defaults.
        """
        opts = {}
        return opts

    def _option_list(self, name, opts):
        """Build the full list of options.

        This includes appending the filenames and incorporating runtime
        options.
        """
        from .base import task_classes, task_type

        deps = self.deps(name)
        options = {}
        options["infile"]    = task_classes["extract"].paths(deps["infile"])[0]
        options["fiberflat"] = task_classes["fiberflatnight"].paths(deps["fiberflat"])[0]
        options["sky"]       = task_classes["sky"].paths(deps["sky"])[0]
        options["models"]    = task_classes["starfit"].paths(deps["models"])[0]
        options["outfile"]   = self.paths(name)[0]

        options.update(opts)
        return option_list(options)

    def _run_cli(self, name, opts, procs, db=None):
        """See BaseTask.run_cli.
        """
        entry = "desi_compute_fluxcalibration"
        optlist = self._option_list(name, opts)
        com = "{} {}".format(entry, " ".join(optlist))
        return com

    def _run(self, name, opts, comm, db=None):
        """See BaseTask.run.
        """
        from ...scripts import fluxcalibration
        optlist = self._option_list(name, opts)
        args = fluxcalibration.parse(optlist)
        fluxcalibration.main(args)
        return

    def postprocessing(self, db, name, cur):
        """For successful runs, postprocessing on DB"""
        # run getready on all fierflatnight with same night,band,spec
        props = self.name_split(name)
        log  = get_logger()
        tt="cframe"
        cmd = "select name from {} where night={} and expid={} and spec={} and band='{}'".format(tt,props["night"],props["expid"],props["spec"],props["band"])
        cur.execute(cmd)
        tasks = [ x for (x,) in cur.fetchall() ]
        log.debug("checking {}".format(tasks))
        for task in tasks :
            task_classes[tt].getready( db=db,name=task,cur=cur)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import (absolute_import, division,
                        print_function, unicode_literals)

__author__ = "Yuji Ikeda"

import numpy as np
from matplotlib.ticker import AutoMinorLocator
from .plotter import Plotter, read_band_labels
from .file_io import read_band_yaml

class BandPlotter(Plotter):
    def load_data(self, data_file="band.yaml"):
        print("Reading band.yaml: ", end="")
        distances, frequencies = read_band_yaml(yaml_file=data_file)
        print("Finished")

        self._distances = distances / distances[-1, -1]
        self._frequencies = frequencies

        band_labels = read_band_labels("band.conf")
        print("band_labels:", band_labels)
        self._band_labels = band_labels

        return self

    def configure(self, ax):
        variables = self._variables

        distances = self._distances

        freq_label = "Frequency ({})".format(variables["freq_unit"])
        d_freq = variables["d_freq"]
        f_min = variables["f_min"]
        f_max = variables["f_max"]
        n_freq = int(round((f_max - f_min) / d_freq)) + 1

        ax.set_xticks([0.0] + list(distances[:, -1]))
        ax.set_xticklabels(self._band_labels)
        ax.set_xlabel("Wave vector")
        ax.set_xlim(distances[0, 0], distances[-1, -1])

        ax.set_yticks(np.linspace(f_min, f_max, n_freq))
        ax.set_ylabel(freq_label)
        ax.set_ylim(f_min, f_max)

        mly = AutoMinorLocator(2)
        ax.yaxis.set_minor_locator(mly)

        for x in [0.0] + list(distances[:, -1]):
            ax.axvline(x, color="k", dashes=(2, 2), linewidth=0.5)
        # for y in np.linspace(f_min, f_max, n_freq):
        #     ax.axhline(y, color="#000000", linestyle=":")
        # zero axis
        ax.axhline(0, color="k", dashes=(2, 2), linewidth=0.5)

    def plot(self, ax):
        variables = self._variables
        distances = self._distances
        frequencies = self._frequencies

        npath, nqpoint, nband = frequencies.shape
        for ipath in range(npath):
            lines = ax.plot(
                distances[ipath],
                frequencies[ipath] * variables["unit"],
                variables["linecolor"],
                dashes=variables["dashes"],
                linewidth=variables["linewidth"],
            )

    def plot_selected_curve(self, ax, ipath, inum):
        variables = self._variables
        distances = self._distances
        frequencies = self._frequencies

        npath, nqpoint, nband = frequencies.shape
        lines = ax.plot(
            distances[ipath],
            frequencies[ipath, :, inum] * variables["unit"],
            variables["linecolor"],
            dashes=variables["dashes"],
            linewidth=variables["linewidth"],
        )

    def create_figure_name(self):
        variables = self._variables
        figure_name = "band_{}.{}".format(
            variables["freq_unit"],
            variables["figure_type"])
        return figure_name
<commit_msg>Modify plotting of grid lines to use ax.grid instead of ax.axvline<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import (absolute_import, division,
                        print_function, unicode_literals)

__author__ = "Yuji Ikeda"

import numpy as np
from matplotlib.ticker import AutoMinorLocator
from .plotter import Plotter, read_band_labels
from .file_io import read_band_yaml

class BandPlotter(Plotter):
    def load_data(self, data_file="band.yaml"):
        print("Reading band.yaml: ", end="")
        distances, frequencies = read_band_yaml(yaml_file=data_file)
        print("Finished")

        self._distances = distances / distances[-1, -1]
        self._frequencies = frequencies

        band_labels = read_band_labels("band.conf")
        print("band_labels:", band_labels)
        self._band_labels = band_labels

        return self

    def configure(self, ax):
        variables = self._variables

        distances = self._distances

        freq_label = "Frequency ({})".format(variables["freq_unit"])
        d_freq = variables["d_freq"]
        f_min = variables["f_min"]
        f_max = variables["f_max"]
        n_freq = int(round((f_max - f_min) / d_freq)) + 1

        ax.set_xticks([0.0] + list(distances[:, -1]))
        ax.set_xticklabels(self._band_labels)
        ax.set_xlabel("Wave vector")
        ax.set_xlim(distances[0, 0], distances[-1, -1])

        ax.set_yticks(np.linspace(f_min, f_max, n_freq))
        ax.set_ylabel(freq_label)
        ax.set_ylim(f_min, f_max)

        mly = AutoMinorLocator(2)
        ax.yaxis.set_minor_locator(mly)

        ax.grid(True, axis='x', color='k', dashes=(2, 2), linewidth=0.5)
        # for y in np.linspace(f_min, f_max, n_freq):
        #     ax.axhline(y, color="#000000", linestyle=":")
        # zero axis
        ax.axhline(0, color="k", dashes=(2, 2), linewidth=0.5)

    def plot(self, ax):
        variables = self._variables
        distances = self._distances
        frequencies = self._frequencies

        npath, nqpoint, nband = frequencies.shape
        for ipath in range(npath):
            lines = ax.plot(
                distances[ipath],
                frequencies[ipath] * variables["unit"],
                variables["linecolor"],
                dashes=variables["dashes"],
                linewidth=variables["linewidth"],
            )

    def plot_selected_curve(self, ax, ipath, inum):
        variables = self._variables
        distances = self._distances
        frequencies = self._frequencies

        npath, nqpoint, nband = frequencies.shape
        lines = ax.plot(
            distances[ipath],
            frequencies[ipath, :, inum] * variables["unit"],
            variables["linecolor"],
            dashes=variables["dashes"],
            linewidth=variables["linewidth"],
        )

    def create_figure_name(self):
        variables = self._variables
        figure_name = "band_{}.{}".format(
            variables["freq_unit"],
            variables["figure_type"])
        return figure_name
<|endoftext|>


<commit_before>import abc
import logging
from datetime import datetime

from modularodm import fields
from modularodm.exceptions import ValidationTypeError, ValidationValueError

from framework.mongo import StoredObject
from website import settings
from website.project.model import User
from website.util import akismet

logger = logging.getLogger(__name__)

def _get_client():
    return akismet.AkismetClient(
        apikey=settings.AKISMET_APIKEY,
        website=settings.DOMAIN,
        verify=True
    )

def _validate_reports(value, *args, **kwargs):
    for key, val in value.iteritems():
        if not User.load(key):
            raise ValidationValueError('Keys must be user IDs')
        if not isinstance(val, dict):
            raise ValidationTypeError('Values must be dictionaries')
        if ('category' not in val or 'text' not in val or 'date' not in val or 'retracted' not in val):
            raise ValidationValueError(
                ('Values must include `date`, `category`, ',
                 '`text`, `retracted` keys')
            )

class SpamStatus(object):
    UNKNOWN = None
    FLAGGED = 1
    SPAM = 2
    HAM = 4

class SpamMixin(StoredObject):
    """Mixin to add to objects that can be marked as spam.
    """

    _meta = {
        'abstract': True
    }

    # # Node fields that trigger an update to search on save
    # SPAM_UPDATE_FIELDS = {
    #     'spam_status',
    # }
    spam_status = fields.IntegerField(default=SpamStatus.UNKNOWN, index=True)
    spam_pro_tip = fields.StringField(default=None)
    # Data representing the original spam indication
    # - author: author name
    # - author_email: email of the author
    # - content: data flagged
    # - headers: request headers
    #   - Remote-Addr: ip address from request
    #   - User-Agent: user agent from request
    #   - Referer: referrer header from request (typo +1, rtd)
    spam_data = fields.DictionaryField(default=dict)
    date_last_reported = fields.DateTimeField(default=None, index=True)

    # Reports is a dict of reports keyed on reporting user
    # Each report is a dictionary including:
    #  - date: date reported
    #  - retracted: if a report has been retracted
    #  - category: What type of spam does the reporter believe this is
    #  - text: Comment on the comment
    reports = fields.DictionaryField(
        default=dict, validate=_validate_reports
    )

    def flag_spam(self):
        # If ham and unedited then tell user that they should read it again
        if self.spam_status == SpamStatus.UNKNOWN:
            self.spam_status = SpamStatus.FLAGGED

    def remove_flag(self, save=False):
        if self.spam_status != SpamStatus.FLAGGED:
            return
        for report in self.reports.values():
            if not report.get('retracted', True):
                return
        self.spam_status = SpamStatus.UNKNOWN
        if save:
            self.save()

    @property
    def is_spam(self):
        return self.spam_status == SpamStatus.SPAM

    @property
    def is_spammy(self):
        return self.spam_status in [SpamStatus.FLAGGED, SpamStatus.SPAM]

    def report_abuse(self, user, save=False, **kwargs):
        """Report object is spam or other abuse of OSF

        :param user: User submitting report
        :param save: Save changes
        :param kwargs: Should include category and message
        :raises ValueError: if user is reporting self
        """
        if user == self.user:
            raise ValueError('User cannot report self.')
        self.flag_spam()
        date = datetime.utcnow()
        report = {'date': date, 'retracted': False}
        report.update(kwargs)
        if 'text' not in report:
            report['text'] = None
        self.reports[user._id] = report
        self.date_last_reported = report['date']
        if save:
            self.save()

    def retract_report(self, user, save=False):
        """Retract last report by user

        Only marks the last report as retracted because there could be
        history in how the object is edited that requires a user
        to flag or retract even if object is marked as HAM.
        :param user: User retracting
        :param save: Save changes
        """
        if user._id in self.reports:
            if not self.reports[user._id]['retracted']:
                self.reports[user._id]['retracted'] = True
                self.remove_flag()
        else:
            raise ValueError('User has not reported this content')
        if save:
            self.save()

    def confirm_ham(self, save=False):
        # not all mixins will implement check spam pre-req, only submit ham when it was incorrectly flagged
        if settings.SPAM_CHECK_ENABLED and self.spam_data and self.spam_status in [SpamStatus.FLAGGED, SpamStatus.SPAM]:
            client = _get_client()
            client.submit_ham(
                user_ip=self.spam_data['headers']['Remote-Addr'],
                user_agent=self.spam_data['headers'].get('User-Agent'),
                referrer=self.spam_data['headers'].get('Referer'),
                comment_content=self.spam_data['content'],
                comment_author=self.spam_data['author'],
                comment_author_email=self.spam_data['author_email'],
            )
            logger.info('confirm_ham update sent')
        self.spam_status = SpamStatus.HAM
        if save:
            self.save()

    def confirm_spam(self, save=False):
        # not all mixins will implement check spam pre-req, only submit spam when it was incorrectly flagged
        if settings.SPAM_CHECK_ENABLED and self.spam_data and self.spam_status == SpamStatus.HAM:
            client = _get_client()
            client.submit_spam(
                user_ip=self.spam_data['headers']['Remote-Addr'],
                user_agent=self.spam_data['headers'].get('User-Agent'),
                referrer=self.spam_data['headers'].get('Referer'),
                comment_content=self.spam_data['content'],
                comment_author=self.spam_data['author'],
                comment_author_email=self.spam_data['author_email'],
            )
            logger.info('confirm_spam update sent')
        self.spam_status = SpamStatus.SPAM
        if save:
            self.save()

    @abc.abstractmethod
    def check_spam(self, saved_fields, request_headers, save=False):
        """Must return is_spam"""
        pass

    def do_check_spam(self, author, author_email, content, request_headers):
        if self.spam_status == SpamStatus.HAM:
            return False
        if self.is_spammy:
            return True

        client = _get_client()
        remote_addr = request_headers['Remote-Addr']
        user_agent = request_headers.get('User-Agent')
        referer = request_headers.get('Referer')
        is_spam, pro_tip = client.check_comment(
            user_ip=remote_addr,
            user_agent=user_agent,
            referrer=referer,
            comment_content=content,
            comment_author=author,
            comment_author_email=author_email
        )
        self.spam_pro_tip = pro_tip
        self.spam_data['headers'] = {
            'Remote-Addr': remote_addr,
            'User-Agent': user_agent,
            'Referer': referer,
        }
        self.spam_data['content'] = content
        self.spam_data['author'] = author
        self.spam_data['author_email'] = author_email
        if is_spam:
            self.flag_spam()
        return is_spam
<commit_msg>spam confirmation should update service on unknown<commit_after>import abc
import logging
from datetime import datetime

from modularodm import fields
from modularodm.exceptions import ValidationTypeError, ValidationValueError

from framework.mongo import StoredObject
from website import settings
from website.project.model import User
from website.util import akismet

logger = logging.getLogger(__name__)

def _get_client():
    return akismet.AkismetClient(
        apikey=settings.AKISMET_APIKEY,
        website=settings.DOMAIN,
        verify=True
    )

def _validate_reports(value, *args, **kwargs):
    for key, val in value.iteritems():
        if not User.load(key):
            raise ValidationValueError('Keys must be user IDs')
        if not isinstance(val, dict):
            raise ValidationTypeError('Values must be dictionaries')
        if ('category' not in val or 'text' not in val or 'date' not in val or 'retracted' not in val):
            raise ValidationValueError(
                ('Values must include `date`, `category`, ',
                 '`text`, `retracted` keys')
            )

class SpamStatus(object):
    UNKNOWN = None
    FLAGGED = 1
    SPAM = 2
    HAM = 4

class SpamMixin(StoredObject):
    """Mixin to add to objects that can be marked as spam.
    """

    _meta = {
        'abstract': True
    }

    # # Node fields that trigger an update to search on save
    # SPAM_UPDATE_FIELDS = {
    #     'spam_status',
    # }
    spam_status = fields.IntegerField(default=SpamStatus.UNKNOWN, index=True)
    spam_pro_tip = fields.StringField(default=None)
    # Data representing the original spam indication
    # - author: author name
    # - author_email: email of the author
    # - content: data flagged
    # - headers: request headers
    #   - Remote-Addr: ip address from request
    #   - User-Agent: user agent from request
    #   - Referer: referrer header from request (typo +1, rtd)
    spam_data = fields.DictionaryField(default=dict)
    date_last_reported = fields.DateTimeField(default=None, index=True)

    # Reports is a dict of reports keyed on reporting user
    # Each report is a dictionary including:
    #  - date: date reported
    #  - retracted: if a report has been retracted
    #  - category: What type of spam does the reporter believe this is
    #  - text: Comment on the comment
    reports = fields.DictionaryField(
        default=dict, validate=_validate_reports
    )

    def flag_spam(self):
        # If ham and unedited then tell user that they should read it again
        if self.spam_status == SpamStatus.UNKNOWN:
            self.spam_status = SpamStatus.FLAGGED

    def remove_flag(self, save=False):
        if self.spam_status != SpamStatus.FLAGGED:
            return
        for report in self.reports.values():
            if not report.get('retracted', True):
                return
        self.spam_status = SpamStatus.UNKNOWN
        if save:
            self.save()

    @property
    def is_spam(self):
        return self.spam_status == SpamStatus.SPAM

    @property
    def is_spammy(self):
        return self.spam_status in [SpamStatus.FLAGGED, SpamStatus.SPAM]

    def report_abuse(self, user, save=False, **kwargs):
        """Report object is spam or other abuse of OSF

        :param user: User submitting report
        :param save: Save changes
        :param kwargs: Should include category and message
        :raises ValueError: if user is reporting self
        """
        if user == self.user:
            raise ValueError('User cannot report self.')
        self.flag_spam()
        date = datetime.utcnow()
        report = {'date': date, 'retracted': False}
        report.update(kwargs)
        if 'text' not in report:
            report['text'] = None
        self.reports[user._id] = report
        self.date_last_reported = report['date']
        if save:
            self.save()

    def retract_report(self, user, save=False):
        """Retract last report by user

        Only marks the last report as retracted because there could be
        history in how the object is edited that requires a user
        to flag or retract even if object is marked as HAM.
        :param user: User retracting
        :param save: Save changes
        """
        if user._id in self.reports:
            if not self.reports[user._id]['retracted']:
                self.reports[user._id]['retracted'] = True
                self.remove_flag()
        else:
            raise ValueError('User has not reported this content')
        if save:
            self.save()

    def confirm_ham(self, save=False):
        # not all mixins will implement check spam pre-req, only submit ham when it was incorrectly flagged
        if settings.SPAM_CHECK_ENABLED and self.spam_data and self.spam_status in [SpamStatus.FLAGGED, SpamStatus.SPAM]:
            client = _get_client()
            client.submit_ham(
                user_ip=self.spam_data['headers']['Remote-Addr'],
                user_agent=self.spam_data['headers'].get('User-Agent'),
                referrer=self.spam_data['headers'].get('Referer'),
                comment_content=self.spam_data['content'],
                comment_author=self.spam_data['author'],
                comment_author_email=self.spam_data['author_email'],
            )
            logger.info('confirm_ham update sent')
        self.spam_status = SpamStatus.HAM
        if save:
            self.save()

    def confirm_spam(self, save=False):
        # not all mixins will implement check spam pre-req, only submit spam when it was incorrectly flagged
        if settings.SPAM_CHECK_ENABLED and self.spam_data and self.spam_status in [SpamStatus.UNKNOWN, SpamStatus.HAM]:
            client = _get_client()
            client.submit_spam(
                user_ip=self.spam_data['headers']['Remote-Addr'],
                user_agent=self.spam_data['headers'].get('User-Agent'),
                referrer=self.spam_data['headers'].get('Referer'),
                comment_content=self.spam_data['content'],
                comment_author=self.spam_data['author'],
                comment_author_email=self.spam_data['author_email'],
            )
            logger.info('confirm_spam update sent')
        self.spam_status = SpamStatus.SPAM
        if save:
            self.save()

    @abc.abstractmethod
    def check_spam(self, saved_fields, request_headers, save=False):
        """Must return is_spam"""
        pass

    def do_check_spam(self, author, author_email, content, request_headers):
        if self.spam_status == SpamStatus.HAM:
            return False
        if self.is_spammy:
            return True

        client = _get_client()
        remote_addr = request_headers['Remote-Addr']
        user_agent = request_headers.get('User-Agent')
        referer = request_headers.get('Referer')
        is_spam, pro_tip = client.check_comment(
            user_ip=remote_addr,
            user_agent=user_agent,
            referrer=referer,
            comment_content=content,
            comment_author=author,
            comment_author_email=author_email
        )
        self.spam_pro_tip = pro_tip
        self.spam_data['headers'] = {
            'Remote-Addr': remote_addr,
            'User-Agent': user_agent,
            'Referer': referer,
        }
        self.spam_data['content'] = content
        self.spam_data['author'] = author
        self.spam_data['author_email'] = author_email
        if is_spam:
            self.flag_spam()
        return is_spam
<|endoftext|>


<commit_before>REPOSITORY_LOCATIONS = dict(
    boringssl = dict(
        # Use commits from branch "chromium-stable-with-bazel"
        commit = "372daf7042ffe3da1335743e7c93d78f1399aba7",  # chromium-68.0.3440.75
        remote = "https://github.com/google/boringssl",
    ),
    com_google_absl = dict(
        commit = "92020a042c0cd46979db9f6f0cb32783dc07765e",  # 2018-06-08
        remote = "https://github.com/abseil/abseil-cpp",
    ),
    com_github_apache_thrift = dict(
        sha256 = "7d59ac4fdcb2c58037ebd4a9da5f9a49e3e034bf75b3f26d9fe48ba3d8806e6b",
        urls = ["https://files.pythonhosted.org/packages/c6/b4/510617906f8e0c5660e7d96fbc5585113f83ad547a3989b80297ac72a74c/thrift-0.11.0.tar.gz"],  # 0.11.0
        strip_prefix = "thrift-0.11.0",
    ),
    com_github_bombela_backward = dict(
        commit = "44ae9609e860e3428cd057f7052e505b4819eb84",  # 2018-02-06
        remote = "https://github.com/bombela/backward-cpp",
    ),
    com_github_circonus_labs_libcircllhist = dict(
        commit = "050da53a44dede7bda136b93a9aeef47bd91fa12",  # 2018-07-02
        remote = "https://github.com/circonus-labs/libcircllhist",
    ),
    com_github_cyan4973_xxhash = dict(
        commit = "7cc9639699f64b750c0b82333dced9ea77e8436e",  # v0.6.5
        remote = "https://github.com/Cyan4973/xxHash",
    ),
    com_github_eile_tclap = dict(
        commit = "3627d9402e529770df9b0edf2aa8c0e0d6c6bb41",  # tclap-1-2-1-release-final
        remote = "https://github.com/eile/tclap",
    ),
    com_github_fmtlib_fmt = dict(
        sha256 = "46628a2f068d0e33c716be0ed9dcae4370242df135aed663a180b9fd8e36733d",
        strip_prefix = "fmt-4.1.0",
        urls = ["https://github.com/fmtlib/fmt/archive/4.1.0.tar.gz"],
    ),
    com_github_gabime_spdlog = dict(
        sha256 = "94f74fd1b3344733d1db3de2ec22e6cbeb769f93a8baa0d4a22b1f62dc7369f8",
        strip_prefix = "spdlog-0.17.0",
        urls = ["https://github.com/gabime/spdlog/archive/v0.17.0.tar.gz"],
    ),
    com_github_gcovr_gcovr = dict(
        commit = "c0d77201039c7b119b18bc7fb991564c602dd75d",
        remote = "https://github.com/gcovr/gcovr",
    ),
    com_github_google_libprotobuf_mutator = dict(
        commit = "c3d2faf04a1070b0b852b0efdef81e1a81ba925e",
        remote = "https://github.com/google/libprotobuf-mutator",
    ),
    com_github_grpc_grpc = dict(
        commit = "bec3b5ada2c5e5d782dff0b7b5018df646b65cb0",  # v1.12.0
        remote = "https://github.com/grpc/grpc.git",
    ),
    io_opentracing_cpp = dict(
        commit = "3b36b084a4d7fffc196eac83203cf24dfb8696b3",  # v1.4.2
        remote = "https://github.com/opentracing/opentracing-cpp",
    ),
    com_lightstep_tracer_cpp = dict(
        commit = "ae6a6bba65f8c4d438a6a3ac855751ca8f52e1dc",
        remote = "https://github.com/lightstep/lightstep-tracer-cpp",  # v0.7.1
    ),
    lightstep_vendored_googleapis = dict(
        commit = "d6f78d948c53f3b400bb46996eb3084359914f9b",
        remote = "https://github.com/google/googleapis",
    ),
    com_github_google_jwt_verify = dict(
        commit = "4eb9e96485b71e00d43acc7207501caafb085b4a",
        remote = "https://github.com/google/jwt_verify_lib",
    ),
    com_github_nodejs_http_parser = dict(
        # 2018-07-20 snapshot to pick up:
        # A performance fix, nodejs/http-parser PR 422.
        # A bug fix, nodejs/http-parser PR 432.
        # TODO(brian-pane): Upgrade to the next http-parser release once it's available
        commit = "77310eeb839c4251c07184a5db8885a572a08352",
        remote = "https://github.com/nodejs/http-parser",
    ),
    com_github_pallets_jinja = dict(
        commit = "78d2f672149e5b9b7d539c575d2c1bfc12db67a9",  # 2.10
        remote = "https://github.com/pallets/jinja",
    ),
    com_github_pallets_markupsafe = dict(
        commit = "d2a40c41dd1930345628ea9412d97e159f828157",  # 1.0
        remote = "https://github.com/pallets/markupsafe",
    ),
    com_github_tencent_rapidjson = dict(
        commit = "f54b0e47a08782a6131cc3d60f94d038fa6e0a51",  # v1.1.0
        remote = "https://github.com/tencent/rapidjson",
    ),
    com_github_twitter_common_lang = dict(
        sha256 = "56d1d266fd4767941d11c27061a57bc1266a3342e551bde3780f9e9eb5ad0ed1",
        urls = ["https://files.pythonhosted.org/packages/08/bc/d6409a813a9dccd4920a6262eb6e5889e90381453a5f58938ba4cf1d9420/twitter.common.lang-0.3.9.tar.gz"],  # 0.3.9
        strip_prefix = "twitter.common.lang-0.3.9/src",
    ),
    com_github_twitter_common_rpc = dict(
        sha256 = "0792b63fb2fb32d970c2e9a409d3d00633190a22eb185145fe3d9067fdaa4514",
        urls = ["https://files.pythonhosted.org/packages/be/97/f5f701b703d0f25fbf148992cd58d55b4d08d3db785aad209255ee67e2d0/twitter.common.rpc-0.3.9.tar.gz"],  # 0.3.9
        strip_prefix = "twitter.common.rpc-0.3.9/src",
    ),
    com_github_twitter_common_finagle_thrift = dict(
        sha256 = "1e3a57d11f94f58745e6b83348ecd4fa74194618704f45444a15bc391fde497a",
        urls = ["https://files.pythonhosted.org/packages/f9/e7/4f80d582578f8489226370762d2cf6bc9381175d1929eba1754e03f70708/twitter.common.finagle-thrift-0.3.9.tar.gz"],  # 0.3.9
        strip_prefix = "twitter.common.finagle-thrift-0.3.9/src",
    ),
    com_google_googletest = dict(
        commit = "43863938377a9ea1399c0596269e0890b5c5515a",
        remote = "https://github.com/google/googletest",
    ),
    com_google_protobuf = dict(
        # TODO(htuch): Switch back to released versions for protobuf when a release > 3.6.0 happens
        # that includes:
        # - https://github.com/google/protobuf/commit/f35669b8d3f46f7f1236bd21f14d744bba251e60
        # - https://github.com/google/protobuf/commit/6a4fec616ec4b20f54d5fb530808b855cb664390
        commit = "6a4fec616ec4b20f54d5fb530808b855cb664390",
        remote = "https://github.com/google/protobuf",
    ),
    grpc_httpjson_transcoding = dict(
        commit = "05a15e4ecd0244a981fdf0348a76658def62fa9c",  # 2018-05-30
        remote = "https://github.com/grpc-ecosystem/grpc-httpjson-transcoding",
    ),
    io_bazel_rules_go = dict(
        commit = "0.11.1",
        remote = "https://github.com/bazelbuild/rules_go",
    ),
    six_archive = dict(
        sha256 = "105f8d68616f8248e24bf0e9372ef04d3cc10104f1980f54d57b2ce73a5ad56a",
        strip_prefix = "",
        urls = ["https://pypi.python.org/packages/source/s/six/six-1.10.0.tar.gz#md5=34eed507548117b2ab523ab14b2f8b55"],
    ),
    # I'd love to name this `com_github_google_subpar`, but something in the Subpar
    # code assumes its repository name is just `subpar`.
    subpar = dict(
        commit = "eb23aa7a5361cabc02464476dd080389340a5522",  # HEAD
        remote = "https://github.com/google/subpar",
    ),
)
<commit_msg>bump abseil-cpp commit (#4034)<commit_after>REPOSITORY_LOCATIONS = dict(
    boringssl = dict(
        # Use commits from branch "chromium-stable-with-bazel"
        commit = "372daf7042ffe3da1335743e7c93d78f1399aba7",  # chromium-68.0.3440.75
        remote = "https://github.com/google/boringssl",
    ),
    com_google_absl = dict(
        commit = "92e07e5590752d6b8e67f7f2f86c6286561e8cea",  # 2018-08-01
        remote = "https://github.com/abseil/abseil-cpp",
    ),
    com_github_apache_thrift = dict(
        sha256 = "7d59ac4fdcb2c58037ebd4a9da5f9a49e3e034bf75b3f26d9fe48ba3d8806e6b",
        urls = ["https://files.pythonhosted.org/packages/c6/b4/510617906f8e0c5660e7d96fbc5585113f83ad547a3989b80297ac72a74c/thrift-0.11.0.tar.gz"],  # 0.11.0
        strip_prefix = "thrift-0.11.0",
    ),
    com_github_bombela_backward = dict(
        commit = "44ae9609e860e3428cd057f7052e505b4819eb84",  # 2018-02-06
        remote = "https://github.com/bombela/backward-cpp",
    ),
    com_github_circonus_labs_libcircllhist = dict(
        commit = "050da53a44dede7bda136b93a9aeef47bd91fa12",  # 2018-07-02
        remote = "https://github.com/circonus-labs/libcircllhist",
    ),
    com_github_cyan4973_xxhash = dict(
        commit = "7cc9639699f64b750c0b82333dced9ea77e8436e",  # v0.6.5
        remote = "https://github.com/Cyan4973/xxHash",
    ),
    com_github_eile_tclap = dict(
        commit = "3627d9402e529770df9b0edf2aa8c0e0d6c6bb41",  # tclap-1-2-1-release-final
        remote = "https://github.com/eile/tclap",
    ),
    com_github_fmtlib_fmt = dict(
        sha256 = "46628a2f068d0e33c716be0ed9dcae4370242df135aed663a180b9fd8e36733d",
        strip_prefix = "fmt-4.1.0",
        urls = ["https://github.com/fmtlib/fmt/archive/4.1.0.tar.gz"],
    ),
    com_github_gabime_spdlog = dict(
        sha256 = "94f74fd1b3344733d1db3de2ec22e6cbeb769f93a8baa0d4a22b1f62dc7369f8",
        strip_prefix = "spdlog-0.17.0",
        urls = ["https://github.com/gabime/spdlog/archive/v0.17.0.tar.gz"],
    ),
    com_github_gcovr_gcovr = dict(
        commit = "c0d77201039c7b119b18bc7fb991564c602dd75d",
        remote = "https://github.com/gcovr/gcovr",
    ),
    com_github_google_libprotobuf_mutator = dict(
        commit = "c3d2faf04a1070b0b852b0efdef81e1a81ba925e",
        remote = "https://github.com/google/libprotobuf-mutator",
    ),
    com_github_grpc_grpc = dict(
        commit = "bec3b5ada2c5e5d782dff0b7b5018df646b65cb0",  # v1.12.0
        remote = "https://github.com/grpc/grpc.git",
    ),
    io_opentracing_cpp = dict(
        commit = "3b36b084a4d7fffc196eac83203cf24dfb8696b3",  # v1.4.2
        remote = "https://github.com/opentracing/opentracing-cpp",
    ),
    com_lightstep_tracer_cpp = dict(
        commit = "ae6a6bba65f8c4d438a6a3ac855751ca8f52e1dc",
        remote = "https://github.com/lightstep/lightstep-tracer-cpp",  # v0.7.1
    ),
    lightstep_vendored_googleapis = dict(
        commit = "d6f78d948c53f3b400bb46996eb3084359914f9b",
        remote = "https://github.com/google/googleapis",
    ),
    com_github_google_jwt_verify = dict(
        commit = "4eb9e96485b71e00d43acc7207501caafb085b4a",
        remote = "https://github.com/google/jwt_verify_lib",
    ),
    com_github_nodejs_http_parser = dict(
        # 2018-07-20 snapshot to pick up:
        # A performance fix, nodejs/http-parser PR 422.
        # A bug fix, nodejs/http-parser PR 432.
        # TODO(brian-pane): Upgrade to the next http-parser release once it's available
        commit = "77310eeb839c4251c07184a5db8885a572a08352",
        remote = "https://github.com/nodejs/http-parser",
    ),
    com_github_pallets_jinja = dict(
        commit = "78d2f672149e5b9b7d539c575d2c1bfc12db67a9",  # 2.10
        remote = "https://github.com/pallets/jinja",
    ),
    com_github_pallets_markupsafe = dict(
        commit = "d2a40c41dd1930345628ea9412d97e159f828157",  # 1.0
        remote = "https://github.com/pallets/markupsafe",
    ),
    com_github_tencent_rapidjson = dict(
        commit = "f54b0e47a08782a6131cc3d60f94d038fa6e0a51",  # v1.1.0
        remote = "https://github.com/tencent/rapidjson",
    ),
    com_github_twitter_common_lang = dict(
        sha256 = "56d1d266fd4767941d11c27061a57bc1266a3342e551bde3780f9e9eb5ad0ed1",
        urls = ["https://files.pythonhosted.org/packages/08/bc/d6409a813a9dccd4920a6262eb6e5889e90381453a5f58938ba4cf1d9420/twitter.common.lang-0.3.9.tar.gz"],  # 0.3.9
        strip_prefix = "twitter.common.lang-0.3.9/src",
    ),
    com_github_twitter_common_rpc = dict(
        sha256 = "0792b63fb2fb32d970c2e9a409d3d00633190a22eb185145fe3d9067fdaa4514",
        urls = ["https://files.pythonhosted.org/packages/be/97/f5f701b703d0f25fbf148992cd58d55b4d08d3db785aad209255ee67e2d0/twitter.common.rpc-0.3.9.tar.gz"],  # 0.3.9
        strip_prefix = "twitter.common.rpc-0.3.9/src",
    ),
    com_github_twitter_common_finagle_thrift = dict(
        sha256 = "1e3a57d11f94f58745e6b83348ecd4fa74194618704f45444a15bc391fde497a",
        urls = ["https://files.pythonhosted.org/packages/f9/e7/4f80d582578f8489226370762d2cf6bc9381175d1929eba1754e03f70708/twitter.common.finagle-thrift-0.3.9.tar.gz"],  # 0.3.9
        strip_prefix = "twitter.common.finagle-thrift-0.3.9/src",
    ),
    com_google_googletest = dict(
        commit = "43863938377a9ea1399c0596269e0890b5c5515a",
        remote = "https://github.com/google/googletest",
    ),
    com_google_protobuf = dict(
        # TODO(htuch): Switch back to released versions for protobuf when a release > 3.6.0 happens
        # that includes:
        # - https://github.com/google/protobuf/commit/f35669b8d3f46f7f1236bd21f14d744bba251e60
        # - https://github.com/google/protobuf/commit/6a4fec616ec4b20f54d5fb530808b855cb664390
        commit = "6a4fec616ec4b20f54d5fb530808b855cb664390",
        remote = "https://github.com/google/protobuf",
    ),
    grpc_httpjson_transcoding = dict(
        commit = "05a15e4ecd0244a981fdf0348a76658def62fa9c",  # 2018-05-30
        remote = "https://github.com/grpc-ecosystem/grpc-httpjson-transcoding",
    ),
    io_bazel_rules_go = dict(
        commit = "0.11.1",
        remote = "https://github.com/bazelbuild/rules_go",
    ),
    six_archive = dict(
        sha256 = "105f8d68616f8248e24bf0e9372ef04d3cc10104f1980f54d57b2ce73a5ad56a",
        strip_prefix = "",
        urls = ["https://pypi.python.org/packages/source/s/six/six-1.10.0.tar.gz#md5=34eed507548117b2ab523ab14b2f8b55"],
    ),
    # I'd love to name this `com_github_google_subpar`, but something in the Subpar
    # code assumes its repository name is just `subpar`.
    subpar = dict(
        commit = "eb23aa7a5361cabc02464476dd080389340a5522",  # HEAD
        remote = "https://github.com/google/subpar",
    ),
)
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
#  restraints_parameterisation.py
#
#  Copyright (C) 2015 Diamond Light Source and STFC Rutherford Appleton
#                     Laboratory, UK.
#
#  Author: David Waterman
#
#  This code is distributed under the BSD license, a copy of which is
#  included in the root directory of this package.

from __future__ import division
from libtbx.phil import parse
from libtbx.utils import Sorry

from dials.algorithms.refinement.restraints.restraints import SingleUnitCellTie

# PHIL options for unit cell restraints
uc_phil_str = '''
restraints
  .help = "Least squares restraints to use in refinement."
  .expert_level = 1
{
  tie_to_target
    .multiple = True
  {
    values = None
      .type = floats(size=6)
      .help = "Target unit cell parameters for the restraint for this"
              "parameterisation"

    sigmas = None
      .help = "The unit cell target values are associated with sigmas which are"
              "used to determine the weight of each restraint. A sigma of zero"
              "or None will remove the restraint at that position"
      .type = floats(size=6, value_min=0.)

    id = None
      .help = "Index of an experiment affected by this restraint. If the"
              "relevant parameterisation affects multiple experiments then any"
              "one of the indices may be supplied"
      .type = int(value_min=0)
  }

  tie_to_group
    .multiple = True
  {
    target = *mean median
      .type = choice
      .help = "Function to tie group parameter values to"

    sigmas = None
      .help = "A sigma of zero or None will remove the restraint at that position"
      .type = floats(value_min=0.)

    id = None
      .help = "Indices of experiments affected by this restraint. For every"
              "parameterisation that requires a restraint at least one"
              "experiment index must be supplied."
      .type = ints(value_min=0)
  }
}

'''

uc_phil_scope = parse(uc_phil_str)

# Define a couple of namedtuple types we will use for convenience
from collections import namedtuple
ParamIndex = namedtuple('ParamIndex', ['parameterisation', 'istart'])
RestraintIndex = namedtuple('RestraintIndex', ['restraint', 'istart'])

class RestraintsParameterisation(object):

  def __init__(self, detector_parameterisations = None,
               beam_parameterisations = None,
               xl_orientation_parameterisations = None,
               xl_unit_cell_parameterisations = None):

    # Keep references to all parameterised models
    self._detector_parameterisations = detector_parameterisations
    self._beam_parameterisations = beam_parameterisations
    self._xl_orientation_parameterisations = xl_orientation_parameterisations
    self._xl_unit_cell_parameterisations = xl_unit_cell_parameterisations

    # Loop over all parameterisations, extract experiment IDs and record
    # global parameter index for each that tells us which parameters have
    # non-zero derivatives
    iparam = 0

    self._exp_to_det_param = {}
    for detp in self._detector_parameterisations:
      for iexp in detp.get_experiment_ids():
        self._exp_to_det_param[iexp] = ParamIndex(detp, iparam)
      iparam += detp.num_free()

    self._exp_to_beam_param = {}
    for beamp in self._beam_parameterisations:
      for iexp in beamp.get_experiment_ids():
        self._exp_to_beam_param[iexp] = ParamIndex(beamp, iparam)
      iparam += beamp.num_free()

    self._exp_to_xlo_param = {}
    for xlop in self._xl_orientation_parameterisations:
      for iexp in xlop.get_experiment_ids():
        self._exp_to_xlo_param[iexp] = ParamIndex(xlop, iparam)
      iparam += xlop.num_free()

    self._exp_to_xluc_param = {}
    for xlucp in self._xl_unit_cell_parameterisations:
      for iexp in xlucp.get_experiment_ids():
        self._exp_to_xluc_param[iexp] = ParamIndex(xlucp, iparam)
      iparam += xlucp.num_free()

    # keep a set that will ensure every model parameterisation only gets
    # a single restraint.
    self._param_to_restraint = set()

    # keep a list of restraint objects that we will add
    self._restraints = []

    return

  #def add_restraints_to_target_detector(self):
  #
  #  return
  #
  #def add_restraints_to_target_beam(self):
  #
  #  return
  #
  #def add_restraints_to_target_xl_orientation(self):
  #
  #  return

  def add_restraints_to_target_xl_unit_cell(self,
    experiment_id, values, sigma):

    # On input we will have one id value, 6 target values and 6 sigmas.

    # select the right parameterisation
    param_i = self._exp_to_xluc_param[experiment_id]

    # fail now if this is already restrained.
    if param_i.parameterisation in self._param_to_restraint:
      raise Sorry("Parameterisation already restrained. Cannot create "
                  "additional restraint with experiment {0}".format(experiment_id))

    # create new restraint
    tie = SingleUnitCellTie(model_parameterisation=param_i.parameterisation,
                            target=values,
                            sigma=sigma)

    # add to the restraint list along with the global parameter index
    self._restraints.append(RestraintIndex(tie, param_i.istart))

    # also add the parameterisation to the set for uniqueness testing
    self._param_to_restraint.add(param_i.parameterisation)

    return

  #def add_restraints_to_group_detector(self):
  #
  #  return
  #
  #def add_restraints_to_group_beam(self):
  #
  #  return
  #
  #def add_restraints_to_group_xl_orientation(self):

    return

  def add_restraints_to_group_xl_unit_cell(self):

    pass
    return

<commit_msg>RestraintsParameterisation.get_values_and_gradients using sparse storage<commit_after>#!/usr/bin/env python
#
#  restraints_parameterisation.py
#
#  Copyright (C) 2015 Diamond Light Source and STFC Rutherford Appleton
#                     Laboratory, UK.
#
#  Author: David Waterman
#
#  This code is distributed under the BSD license, a copy of which is
#  included in the root directory of this package.

from __future__ import division
from libtbx.phil import parse
from libtbx.utils import Sorry
from scitbx.array_family import flex
from scitbx import sparse

from dials.algorithms.refinement.restraints.restraints import SingleUnitCellTie

# PHIL options for unit cell restraints
uc_phil_str = '''
restraints
  .help = "Least squares restraints to use in refinement."
  .expert_level = 1
{
  tie_to_target
    .multiple = True
  {
    values = None
      .type = floats(size=6)
      .help = "Target unit cell parameters for the restraint for this"
              "parameterisation"

    sigmas = None
      .help = "The unit cell target values are associated with sigmas which are"
              "used to determine the weight of each restraint. A sigma of zero"
              "or None will remove the restraint at that position"
      .type = floats(size=6, value_min=0.)

    id = None
      .help = "Index of an experiment affected by this restraint. If the"
              "relevant parameterisation affects multiple experiments then any"
              "one of the indices may be supplied"
      .type = int(value_min=0)
  }

  tie_to_group
    .multiple = True
  {
    target = *mean median
      .type = choice
      .help = "Function to tie group parameter values to"

    sigmas = None
      .help = "A sigma of zero or None will remove the restraint at that position"
      .type = floats(value_min=0.)

    id = None
      .help = "Indices of experiments affected by this restraint. For every"
              "parameterisation that requires a restraint at least one"
              "experiment index must be supplied."
      .type = ints(value_min=0)
  }
}

'''

uc_phil_scope = parse(uc_phil_str)

# Define a couple of namedtuple types we will use for convenience
from collections import namedtuple
ParamIndex = namedtuple('ParamIndex', ['parameterisation', 'istart'])
RestraintIndex = namedtuple('RestraintIndex', ['restraint', 'istart'])

class RestraintsParameterisation(object):

  def __init__(self, detector_parameterisations = None,
               beam_parameterisations = None,
               xl_orientation_parameterisations = None,
               xl_unit_cell_parameterisations = None):

    # Keep references to all parameterised models
    self._detector_parameterisations = detector_parameterisations
    self._beam_parameterisations = beam_parameterisations
    self._xl_orientation_parameterisations = xl_orientation_parameterisations
    self._xl_unit_cell_parameterisations = xl_unit_cell_parameterisations

    # Loop over all parameterisations, extract experiment IDs and record
    # global parameter index for each that tells us which parameters have
    # non-zero derivatives
    iparam = 0

    self._exp_to_det_param = {}
    for detp in self._detector_parameterisations:
      for iexp in detp.get_experiment_ids():
        self._exp_to_det_param[iexp] = ParamIndex(detp, iparam)
      iparam += detp.num_free()

    self._exp_to_beam_param = {}
    for beamp in self._beam_parameterisations:
      for iexp in beamp.get_experiment_ids():
        self._exp_to_beam_param[iexp] = ParamIndex(beamp, iparam)
      iparam += beamp.num_free()

    self._exp_to_xlo_param = {}
    for xlop in self._xl_orientation_parameterisations:
      for iexp in xlop.get_experiment_ids():
        self._exp_to_xlo_param[iexp] = ParamIndex(xlop, iparam)
      iparam += xlop.num_free()

    self._exp_to_xluc_param = {}
    for xlucp in self._xl_unit_cell_parameterisations:
      for iexp in xlucp.get_experiment_ids():
        self._exp_to_xluc_param[iexp] = ParamIndex(xlucp, iparam)
      iparam += xlucp.num_free()

    # the number of free parameters
    self._nparam = iparam

    # keep a set that will ensure every model parameterisation only gets
    # a single restraint.
    self._param_to_restraint = set()

    # keep a list of restraint objects that we will add
    self._restraints = []

    return

  #def add_restraints_to_target_detector(self):
  #
  #  return
  #
  #def add_restraints_to_target_beam(self):
  #
  #  return
  #
  #def add_restraints_to_target_xl_orientation(self):
  #
  #  return

  def add_restraints_to_target_xl_unit_cell(self,
    experiment_id, values, sigma):

    # On input we will have one id value, 6 target values and 6 sigmas.

    # select the right parameterisation
    param_i = self._exp_to_xluc_param[experiment_id]

    # fail now if this is already restrained.
    if param_i.parameterisation in self._param_to_restraint:
      raise Sorry("Parameterisation already restrained. Cannot create "
                  "additional restraint with experiment {0}".format(experiment_id))

    # create new restraint
    tie = SingleUnitCellTie(model_parameterisation=param_i.parameterisation,
                            target=values,
                            sigma=sigma)

    # add to the restraint list along with the global parameter index
    self._restraints.append(RestraintIndex(tie, param_i.istart))

    # also add the parameterisation to the set for uniqueness testing
    self._param_to_restraint.add(param_i.parameterisation)

    return

  #def add_restraints_to_group_detector(self):
  #
  #  return
  #
  #def add_restraints_to_group_beam(self):
  #
  #  return
  #
  #def add_restraints_to_group_xl_orientation(self):
  #
  #  return

  def add_restraints_to_group_xl_unit_cell(self):

    pass
    return

  def get_values_and_gradients(self):

    values = []
    for r in self._restraints:
      values.extend(r.restraint.values())

    nrows = len(values)
    gradients = sparse.matrix(nrows, self._nparam)

    for irow, r in enumerate(self._restraints):
      icol = r.istart
      # convert square list-of-lists into a 2D array for block assignment
      grads = flex.double(r.restraint.gradients())
      gradients.assign_block(grads, irow, icol)

    return values, gradients
<|endoftext|>