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

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

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

from db import db

LOGIN_HARDENING = True
# Unfortunately, the login hardening measures mess with the tests in
# non-deterministic ways.  TODO rewrite the tests so we can more
# precisely control which code paths are exercised.
if os.environ.get('SECUREDROP_ENV') == 'test':
    LOGIN_HARDENING = False

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

class InvalidUsernameException(Exception):

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

class LoginThrottledException(Exception):

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

class WrongPasswordException(Exception):

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

class BadTokenException(Exception):

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

class PasswordError(Exception):

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

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

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

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

class NonDicewarePassword(PasswordError):

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

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

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

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

    MIN_USERNAME_LEN = 3

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

        if otp_secret:
            self.set_hotp_secret(otp_secret)

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

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

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

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

    MAX_PASSWORD_LEN = 128
    MIN_PASSWORD_LEN = 14

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    _LOGIN_ATTEMPT_PERIOD = 60  # seconds
    _MAX_LOGIN_ATTEMPTS_PER_PERIOD = 5

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

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

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

        if LOGIN_HARDENING:
            cls.throttle_login(user)

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

class JournalistLoginAttempt(db.Model):

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

    def __init__(self, journalist):
        self.journalist_id = journalist.id
<commit_msg>Journalist interface: Login throttling should be per-user<commit_after># -*- coding: utf-8 -*-
import binascii
import datetime
import base64
import os
import scrypt
import pyotp
import qrcode
# Using svg because it doesn't require additional dependencies
import qrcode.image.svg

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

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

from db import db

LOGIN_HARDENING = True
# Unfortunately, the login hardening measures mess with the tests in
# non-deterministic ways.  TODO rewrite the tests so we can more
# precisely control which code paths are exercised.
if os.environ.get('SECUREDROP_ENV') == 'test':
    LOGIN_HARDENING = False

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

class InvalidUsernameException(Exception):

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

class LoginThrottledException(Exception):

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

class WrongPasswordException(Exception):

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

class BadTokenException(Exception):

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

class PasswordError(Exception):

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

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

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

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

class NonDicewarePassword(PasswordError):

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

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

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

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

    MIN_USERNAME_LEN = 3

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

        if otp_secret:
            self.set_hotp_secret(otp_secret)

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

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

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

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

    MAX_PASSWORD_LEN = 128
    MIN_PASSWORD_LEN = 14

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    _LOGIN_ATTEMPT_PERIOD = 60  # seconds
    _MAX_LOGIN_ATTEMPTS_PER_PERIOD = 5

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

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

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

        if LOGIN_HARDENING:
            cls.throttle_login(user)

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

class JournalistLoginAttempt(db.Model):

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

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


<commit_before>#!/usr/bin/env python3

import pathlib

input_file = pathlib.Path(__file__).parent.parent.joinpath('day18_input.txt')

class Grid:
    """An animated light grid"""
    @staticmethod
    def _parse_input(text):
        return [1 if s == '#' else 0 for s in text.replace('\n', '')]

    def __init__(self, lights_string, width=None, broken_corners=False):
        """Create grid from the puzzle input (lights_string)

        Raises ValueError if width doesn't divide cleanly into total length,
        or if the grid is not a square and the width is not provided.
        """
        parsed = self._parse_input(lights_string)
        self.lights = parsed
        self.total_lights = len(self.lights)

        if width is None:
            width = self.total_lights ** 0.5
            if not width.is_integer():
                raise ValueError(
                    f'Grid is not a square and width was not provided.')
        self.width = int(width)

        if self.total_lights % self.width:
            raise ValueError(
                f'Grid of length {len(lights_string)} and width {width}'
                ' is not rectangular.')

        self.broken_corners = broken_corners
        if broken_corners:
            self.corner_indices = (
                0, self.width - 1,
                self.total_lights - self.width,
                self.total_lights - 1)
            for idx in self.corner_indices:
                self.lights[idx] = 1

    def _neighbour_indices(self, index):
        """Return index’s neighbours on a rectangular grid

        The 'grid' is virtual so the specified width and
        length of the light list are used to determine if
        positions are on the edge or in the corners, and
        only return indices for neighbours inside the grid.
        """
        w = self.width

        # These aren't always safe, but filtered later
        # along with the upper and lower diagonals.
        positions = [-w, w]
        if index % w != 0:  # Not on the left-hand edge
            positions.extend([-(w + 1), -1, w - 1])
        if (index + 1) % w != 0:  # Not on the right-hand edge
            positions.extend([-(w - 1), 1, w + 1])

        indices = [i for i in (index + p for p in positions)
                   if 0 <= i < self.total_lights]
        return indices

    def neighbours_lit(self, target, light_state):
        """Return the number of neighbours that are turned on

        light_state is a list of lights, so toggles don't affect
        other toggles in the same animation step.
        """
        indices = self._neighbour_indices(target)
        return sum(light_state[idx] for idx in indices)

    def toggle(self, index, light_state):
        """Toggle the light at index based on neighbours

        light_state is a list of lights, so toggles don't affect
        other toggles in the same animation step.

        A light which is on stays on when 2 or 3 neighbours
        are on, and turns off otherwise.

        A light which is off turns on if exactly 3 neighbours
        are on, and stays off otherwise.
        """
        if self.broken_corners and index in self.corner_indices:
            return
        lit = light_state[index]
        neighbour_score = self.neighbours_lit(index, light_state)
        if lit and neighbour_score not in (2, 3):
            self.lights[index] = 0
        elif not lit and neighbour_score == 3:
            self.lights[index] = 1

    def animate(self, to_stage):
        """Animate the grid by performing repeated rounds of toggles"""
        for stage in range(to_stage):
            light_state = self.lights[:]
            for idx in range(self.total_lights):
                self.toggle(idx, light_state)
        return self.lights[:]

def test_parse():
    assert Grid._parse_input('......\n######') == [0, 0, 0, 0, 0, 0,
                                                   1, 1, 1, 1, 1, 1]

def test_misshaped_grid():
    try:
        Grid(lights_string=('.' * 5), width=2)
    except ValueError:
        assert True
    else:
        assert False

def test_neighbours():
    g = Grid(lights_string=('#' * 36), width=6)
    assert sorted(g._neighbour_indices(0)) == [1, 6, 7]
    assert sorted(g._neighbour_indices(1)) == [0, 2, 6, 7, 8]
    assert sorted(g._neighbour_indices(5)) == [4, 10, 11]
    assert sorted(g._neighbour_indices(6)) == [0, 1, 7, 12, 13]
    assert sorted(g._neighbour_indices(7)) == [0, 1, 2, 6, 8, 12, 13, 14]
    assert sorted(g._neighbour_indices(29)) == [22, 23, 28, 34, 35]
    assert sorted(g._neighbour_indices(35)) == [28, 29, 34]

def test_neighbours_lit():
    lights = '''\
....
.##.
.##.
....'''
    g = Grid(lights_string=lights, width=4)
    light_state = g.lights[:]
    assert g.neighbours_lit(0, light_state) == 1
    assert g.neighbours_lit(4, light_state) == 2
    assert g.neighbours_lit(5, light_state) == 3
    assert g.neighbours_lit(14, light_state) == 2

def test_toggle():
    lights = '''\
....
.###
.##.
....'''
    g = Grid(lights_string=lights, width=4)
    light_state = g.lights[:]
    for index, expected in [(0, 0), (2, 1), (11, 1), (14, 0)]:
        g.toggle(index, light_state)
        assert g.lights[index] == expected

def test_sample():
    lights = '''\
.#.#.#
...##.
#....#
..#...
#.#..#
####..'''
    g = Grid(lights_string=lights)
    assert sum(g.animate(5)) == 4
    bg = Grid(lights_string=lights, broken_corners=True)
    assert sum(bg.animate(5)) == 17

def main():
    grid_text = input_file.read_text()
    grid = Grid(lights_string=grid_text, width=100)
    print(sum(grid.animate(100)))
    broken_grid = Grid(lights_string=grid_text, width=100,
                       broken_corners=True)
    print(sum(broken_grid.animate(100)))

if __name__ == '__main__':
    main()
<commit_msg>Day 18: Add str method to Grid (to print the game of life!)<commit_after>#!/usr/bin/env python3

import pathlib

input_file = pathlib.Path(__file__).parent.parent.joinpath('day18_input.txt')

class Grid:
    """An animated light grid"""
    @staticmethod
    def _parse_input(text):
        return [1 if s == '#' else 0 for s in text.replace('\n', '')]

    def __init__(self, lights_string, width=None, broken_corners=False):
        """Create grid from the puzzle input (lights_string)

        Raises ValueError if width doesn't divide cleanly into total length,
        or if the grid is not a square and the width is not provided.
        """
        parsed = self._parse_input(lights_string)
        self.lights = parsed
        self.total_lights = len(self.lights)

        if width is None:
            width = self.total_lights ** 0.5
            if not width.is_integer():
                raise ValueError(
                    f'Grid is not a square and width was not provided.')
        self.width = int(width)

        if self.total_lights % self.width:
            raise ValueError(
                f'Grid of length {len(lights_string)} and width {width}'
                ' is not rectangular.')

        self.broken_corners = broken_corners
        if broken_corners:
            self.corner_indices = (
                0, self.width - 1,
                self.total_lights - self.width,
                self.total_lights - 1)
            for idx in self.corner_indices:
                self.lights[idx] = 1

    def __str__(self):
        rows = (
            slice(start, end)
            for start, end in zip(
                range(0, self.total_lights, self.width),
                range(self.width, self.total_lights + self.width, self.width)
            ))
        result = []
        for row_slice in rows:
            result.append(
                ''.join(['#' if l else '.' for l in self.lights[row_slice]]))
        return '\n'.join(result)

    def _neighbour_indices(self, index):
        """Return index’s neighbours on a rectangular grid

        The 'grid' is virtual so the specified width and
        length of the light list are used to determine if
        positions are on the edge or in the corners, and
        only return indices for neighbours inside the grid.
        """
        w = self.width

        # These aren't always safe, but filtered later
        # along with the upper and lower diagonals.
        positions = [-w, w]
        if index % w != 0:  # Not on the left-hand edge
            positions.extend([-(w + 1), -1, w - 1])
        if (index + 1) % w != 0:  # Not on the right-hand edge
            positions.extend([-(w - 1), 1, w + 1])

        indices = [i for i in (index + p for p in positions)
                   if 0 <= i < self.total_lights]
        return indices

    def neighbours_lit(self, target, light_state):
        """Return the number of neighbours that are turned on

        light_state is a list of lights, so toggles don't affect
        other toggles in the same animation step.
        """
        indices = self._neighbour_indices(target)
        return sum(light_state[idx] for idx in indices)

    def toggle(self, index, light_state):
        """Toggle the light at index based on neighbours

        light_state is a list of lights, so toggles don't affect
        other toggles in the same animation step.

        A light which is on stays on when 2 or 3 neighbours
        are on, and turns off otherwise.

        A light which is off turns on if exactly 3 neighbours
        are on, and stays off otherwise.
        """
        if self.broken_corners and index in self.corner_indices:
            return
        lit = light_state[index]
        neighbour_score = self.neighbours_lit(index, light_state)
        if lit and neighbour_score not in (2, 3):
            self.lights[index] = 0
        elif not lit and neighbour_score == 3:
            self.lights[index] = 1

    def animate(self, to_stage):
        """Animate the grid by performing repeated rounds of toggles"""
        for stage in range(to_stage):
            light_state = self.lights[:]
            for idx in range(self.total_lights):
                self.toggle(idx, light_state)
        return self.lights[:]

def test_parse():
    assert Grid._parse_input('......\n######') == [0, 0, 0, 0, 0, 0,
                                                   1, 1, 1, 1, 1, 1]

def test_misshaped_grid():
    try:
        Grid(lights_string=('.' * 5), width=2)
    except ValueError:
        assert True
    else:
        assert False

def test_neighbours():
    g = Grid(lights_string=('#' * 36), width=6)
    assert sorted(g._neighbour_indices(0)) == [1, 6, 7]
    assert sorted(g._neighbour_indices(1)) == [0, 2, 6, 7, 8]
    assert sorted(g._neighbour_indices(5)) == [4, 10, 11]
    assert sorted(g._neighbour_indices(6)) == [0, 1, 7, 12, 13]
    assert sorted(g._neighbour_indices(7)) == [0, 1, 2, 6, 8, 12, 13, 14]
    assert sorted(g._neighbour_indices(29)) == [22, 23, 28, 34, 35]
    assert sorted(g._neighbour_indices(35)) == [28, 29, 34]

def test_neighbours_lit():
    lights = '''\
....
.##.
.##.
....'''
    g = Grid(lights_string=lights, width=4)
    light_state = g.lights[:]
    assert g.neighbours_lit(0, light_state) == 1
    assert g.neighbours_lit(4, light_state) == 2
    assert g.neighbours_lit(5, light_state) == 3
    assert g.neighbours_lit(14, light_state) == 2

def test_toggle():
    lights = '''\
....
.###
.##.
....'''
    g = Grid(lights_string=lights, width=4)
    light_state = g.lights[:]
    for index, expected in [(0, 0), (2, 1), (11, 1), (14, 0)]:
        g.toggle(index, light_state)
        assert g.lights[index] == expected

def test_sample():
    lights = '''\
.#.#.#
...##.
#....#
..#...
#.#..#
####..'''
    g = Grid(lights_string=lights)
    assert sum(g.animate(5)) == 4
    bg = Grid(lights_string=lights, broken_corners=True)
    assert sum(bg.animate(5)) == 17

def main():
    grid_text = input_file.read_text()
    grid = Grid(lights_string=grid_text, width=100)
    print(sum(grid.animate(100)))
    broken_grid = Grid(lights_string=grid_text, width=100,
                       broken_corners=True)
    print(sum(broken_grid.animate(100)))

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before># Copyright 2017-2020 TensorHub, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division

import ast
import json
import logging
import os
import re

import six

from guild import config as configlib
from guild import flag_util
from guild import guildfile
from guild import model as modellib
from guild import plugin as pluginlib
from guild import python_util
from guild import util

log = logging.getLogger("guild")

class NotebookPlugin(pluginlib.Plugin):

    resolve_model_op_priority = 60

    def resolve_model_op(self, opspec):
        path = os.path.join(configlib.cwd(), opspec)
        if _is_notebook(path):
            model = NotebookModelProxy(path, opspec)
            return model, model.op_name
        return None

    def guildfile_loaded(self, gf):
        cache = {}
        for m in gf.models.values():
            for opdef in m.operations:
                if _is_nbexec_op(opdef):
                    _apply_notebook_flags(opdef, cache)

def _is_notebook(path):
    return os.path.isfile(path) and path.lower().endswith(".ipynb")

def _is_nbexec_op(opdef):
    return opdef.main and opdef.main.rstrip().startswith("guild.plugins.nbexec")

def _apply_notebook_flags(opdef, cache):
    notebook_path = _nbexec_notebook_path(opdef)
    if not notebook_path or not os.path.exists(notebook_path):
        return
    flags_data = _flags_data_for_notebook(notebook_path, opdef, cache)
    flag_util.apply_flags_data_to_op(flags_data, opdef)

def _nbexec_notebook_path(opdef):
    argv = util.shlex_split(opdef.main)
    assert argv and argv[0] == "guild.plugins.nbexec", (opdef.main, argv)
    if len(argv) > 1:
        return os.path.join(opdef.guildfile.dir, argv[1])
    return None

class NotebookModelProxy(object):

    name = ""

    def __init__(self, notebook_path, op_name):
        self.notebook_path = notebook_path
        if os.path.isabs(op_name) or op_name.startswith(".."):
            self.op_name = os.path.basename(op_name)
        else:
            self.op_name = op_name
        self.modeldef = _init_modeldef(self.notebook_path, self.name, self.op_name)
        script_base = notebook_path[: -len(self.op_name)]
        self.reference = modellib.script_model_ref(self.name, script_base)

def _init_modeldef(notebook_path, model_name, op_name):
    data = [
        {
            "model": model_name,
            "operations": {
                op_name: {
                    "main": "guild.plugins.nbexec %s" % _normpath_for_main(notebook_path),
                    "flags": _flags_data_for_notebook(notebook_path),
                }
            },
        }
    ]
    gf = guildfile.Guildfile(data, dir=os.path.dirname(notebook_path))
    return gf.models[model_name]

def _normpath_for_main(path):
    """Normalize path for use in operation main spec.

    Guild does not support Windows style backslashes in op main specs
    so we normalize to forward slash.
    """
    return path.replace(os.path.sep, "/")

def _flags_data_for_notebook(notebook_path, opdef=None, cache=None):
    cache = cache or {}
    try:
        return cache[notebook_path]
    except KeyError:
        data = dict(_iter_notebook_flag_data(notebook_path, opdef))
        return cache.setdefault(notebook_path, data)

def _iter_notebook_flag_data(notebook_path, opdef=None):
    """Iterator that returns name, flag_data for notebooks flags.

    Notebook flag data is derrived from inspecting the notebook for
    global assigns and by applying replacement patterns optionally
    provided by operation flag defs.

    If opdef is provided and a flag provides a replacement patter, it
    takes precedence over values from global assigns.
    """
    seen = set()
    if opdef:
        for name, val in _iter_notebook_flag_replacements(notebook_path, opdef):
            yield name, _flag_data_for_val(val)
            seen.add(name)
    for name, val in _iter_notebook_assigns(notebook_path):
        if not name in seen:
            yield name, _flag_data_for_val(val)

def _iter_notebook_flag_replacements(notebook_path, opdef):
    nb_data = _load_notebook(notebook_path)
    if not nb_data:
        return
    for src in _iter_notebook_source(nb_data):
        for name, val in _iter_replacements_for_source(src, opdef):
            yield name, val

def _iter_replacements_for_source(src, opdef):
    for flagdef in opdef.flags:
        nb_replace = flagdef.extra.get("nb-replace")
        if not nb_replace:
            continue
        try:
            m = re.search(nb_replace, src, re.MULTILINE)
        except ValueError:
            pass
        else:
            if m and m.regs:
                val_str = _replacement_capture(m, src)
                try:
                    val = eval(val_str)
                except Exception as e:
                    if log.getEffectiveLevel() < logging.DEBUG:
                        log.exception("evaluating captured string %r", val_str)
                    else:
                        log.warning("cannot decode captured string %r: %s", val_str, e)
                else:
                    yield flagdef.name, val

def _replacement_capture(m, s):
    start_slice, end_slice = _replacement_capture_slice(m)
    return s[start_slice:end_slice]

def _replacement_capture_slice(m):
    if len(m.regs) == 1:
        # If no capture groups, assume entire capture.
        return m.regs[0]
    else:
        # Use the first capture group.
        return m.regs[1]

def _iter_notebook_assigns(notebook_path):
    nb_data = _load_notebook(notebook_path)
    if not nb_data:
        return
    for src in _iter_notebook_source(nb_data):
        try:
            src = _ipython_to_python(src)
        except MissingIPython:
            break
        else:
            for _assign_node, target_node, _val_node, val in _iter_source_val_assigns(
                src
            ):
                yield target_node.id, val

class MissingIPython(Exception):
    pass

def _ipython_to_python(source):
    try:
        from IPython.core.inputsplitter import IPythonInputSplitter
    except ImportError:
        log.warning(
            "IPython is required to process Notebook source code - "
            "install it by running 'pip install ipython'"
        )
        raise MissingIPython()
    else:
        isp = IPythonInputSplitter(line_input_checker=False)
        python_source = isp.transform_cell(source)
        return _fix_ipython_to_python_lf(python_source, source)

def _fix_ipython_to_python_lf(ipython_s, python_s):
    if python_s[-1:] != "\n" and ipython_s[-1:] == "\n":
        return ipython_s[:-1]
    return ipython_s

def _load_notebook(path):
    try:
        return json.load(open(path))
    except Exception as e:
        if log.getEffectiveLevel() <= logging.DEBUG:
            log.exception("load notebook '%s'", path)
        else:
            log.warning("error loading notebook '%s': %s", path, e)
            return None

def _iter_notebook_source(nb_data):
    for cell in nb_data.get("cells", []):
        if cell.get("cell_type") == "code":
            yield "".join(cell.get("source", []))

def _iter_source_val_assigns(src):
    for node in _iter_source_assigns(src):
        try:
            val = python_util.ast_param_val(node.value)
        except TypeError:
            pass
        else:
            # Only support assigns to right-most name targets.
            target = node.targets[-1]
            if isinstance(target, ast.Name):
                yield node, target, node.value, val

def _iter_source_assigns(s):
    try:
        m = ast.parse(s)
    except Exception as e:
        if log.getEffectiveLevel() <= logging.DEBUG:
            log.exception("parsing %r", s)
        else:
            log.warning("error parsing Python source %r: %s", s, e)
    else:
        for node in m.body:
            if isinstance(node, ast.Assign):
                yield node

def _flag_data_for_val(val):
    return {
        "default": val,
        "type": _flag_type(val),
    }

def _flag_type(val):
    if isinstance(val, six.string_types):
        return "string"
    elif isinstance(val, bool):
        return "boolean"
    elif isinstance(val, (int, float)):
        return "number"
    else:
        return None
<commit_msg>Whitespace<commit_after># Copyright 2017-2020 TensorHub, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division

import ast
import json
import logging
import os
import re

import six

from guild import config as configlib
from guild import flag_util
from guild import guildfile
from guild import model as modellib
from guild import plugin as pluginlib
from guild import python_util
from guild import util

log = logging.getLogger("guild")

class NotebookPlugin(pluginlib.Plugin):

    resolve_model_op_priority = 60

    def resolve_model_op(self, opspec):
        path = os.path.join(configlib.cwd(), opspec)
        if _is_notebook(path):
            model = NotebookModelProxy(path, opspec)
            return model, model.op_name
        return None

    def guildfile_loaded(self, gf):
        cache = {}
        for m in gf.models.values():
            for opdef in m.operations:
                if _is_nbexec_op(opdef):
                    _apply_notebook_flags(opdef, cache)

def _is_notebook(path):
    return os.path.isfile(path) and path.lower().endswith(".ipynb")

def _is_nbexec_op(opdef):
    return opdef.main and opdef.main.rstrip().startswith("guild.plugins.nbexec")

def _apply_notebook_flags(opdef, cache):
    notebook_path = _nbexec_notebook_path(opdef)
    if not notebook_path or not os.path.exists(notebook_path):
        return
    flags_data = _flags_data_for_notebook(notebook_path, opdef, cache)
    flag_util.apply_flags_data_to_op(flags_data, opdef)

def _nbexec_notebook_path(opdef):
    argv = util.shlex_split(opdef.main)
    assert argv and argv[0] == "guild.plugins.nbexec", (opdef.main, argv)
    if len(argv) > 1:
        return os.path.join(opdef.guildfile.dir, argv[1])
    return None

class NotebookModelProxy(object):

    name = ""

    def __init__(self, notebook_path, op_name):
        self.notebook_path = notebook_path
        if os.path.isabs(op_name) or op_name.startswith(".."):
            self.op_name = os.path.basename(op_name)
        else:
            self.op_name = op_name
        self.modeldef = _init_modeldef(self.notebook_path, self.name, self.op_name)
        script_base = notebook_path[: -len(self.op_name)]
        self.reference = modellib.script_model_ref(self.name, script_base)

def _init_modeldef(notebook_path, model_name, op_name):
    data = [
        {
            "model": model_name,
            "operations": {
                op_name: {
                    "main": "guild.plugins.nbexec %s"
                    % _normpath_for_main(notebook_path),
                    "flags": _flags_data_for_notebook(notebook_path),
                }
            },
        }
    ]
    gf = guildfile.Guildfile(data, dir=os.path.dirname(notebook_path))
    return gf.models[model_name]

def _normpath_for_main(path):
    """Normalize path for use in operation main spec.

    Guild does not support Windows style backslashes in op main specs
    so we normalize to forward slash.
    """
    return path.replace(os.path.sep, "/")

def _flags_data_for_notebook(notebook_path, opdef=None, cache=None):
    cache = cache or {}
    try:
        return cache[notebook_path]
    except KeyError:
        data = dict(_iter_notebook_flag_data(notebook_path, opdef))
        return cache.setdefault(notebook_path, data)

def _iter_notebook_flag_data(notebook_path, opdef=None):
    """Iterator that returns name, flag_data for notebooks flags.

    Notebook flag data is derrived from inspecting the notebook for
    global assigns and by applying replacement patterns optionally
    provided by operation flag defs.

    If opdef is provided and a flag provides a replacement patter, it
    takes precedence over values from global assigns.
    """
    seen = set()
    if opdef:
        for name, val in _iter_notebook_flag_replacements(notebook_path, opdef):
            yield name, _flag_data_for_val(val)
            seen.add(name)
    for name, val in _iter_notebook_assigns(notebook_path):
        if not name in seen:
            yield name, _flag_data_for_val(val)

def _iter_notebook_flag_replacements(notebook_path, opdef):
    nb_data = _load_notebook(notebook_path)
    if not nb_data:
        return
    for src in _iter_notebook_source(nb_data):
        for name, val in _iter_replacements_for_source(src, opdef):
            yield name, val

def _iter_replacements_for_source(src, opdef):
    for flagdef in opdef.flags:
        nb_replace = flagdef.extra.get("nb-replace")
        if not nb_replace:
            continue
        try:
            m = re.search(nb_replace, src, re.MULTILINE)
        except ValueError:
            pass
        else:
            if m and m.regs:
                val_str = _replacement_capture(m, src)
                try:
                    val = eval(val_str)
                except Exception as e:
                    if log.getEffectiveLevel() < logging.DEBUG:
                        log.exception("evaluating captured string %r", val_str)
                    else:
                        log.warning("cannot decode captured string %r: %s", val_str, e)
                else:
                    yield flagdef.name, val

def _replacement_capture(m, s):
    start_slice, end_slice = _replacement_capture_slice(m)
    return s[start_slice:end_slice]

def _replacement_capture_slice(m):
    if len(m.regs) == 1:
        # If no capture groups, assume entire capture.
        return m.regs[0]
    else:
        # Use the first capture group.
        return m.regs[1]

def _iter_notebook_assigns(notebook_path):
    nb_data = _load_notebook(notebook_path)
    if not nb_data:
        return
    for src in _iter_notebook_source(nb_data):
        try:
            src = _ipython_to_python(src)
        except MissingIPython:
            break
        else:
            for _assign_node, target_node, _val_node, val in _iter_source_val_assigns(
                src
            ):
                yield target_node.id, val

class MissingIPython(Exception):
    pass

def _ipython_to_python(source):
    try:
        from IPython.core.inputsplitter import IPythonInputSplitter
    except ImportError:
        log.warning(
            "IPython is required to process Notebook source code - "
            "install it by running 'pip install ipython'"
        )
        raise MissingIPython()
    else:
        isp = IPythonInputSplitter(line_input_checker=False)
        python_source = isp.transform_cell(source)
        return _fix_ipython_to_python_lf(python_source, source)

def _fix_ipython_to_python_lf(ipython_s, python_s):
    if python_s[-1:] != "\n" and ipython_s[-1:] == "\n":
        return ipython_s[:-1]
    return ipython_s

def _load_notebook(path):
    try:
        return json.load(open(path))
    except Exception as e:
        if log.getEffectiveLevel() <= logging.DEBUG:
            log.exception("load notebook '%s'", path)
        else:
            log.warning("error loading notebook '%s': %s", path, e)
            return None

def _iter_notebook_source(nb_data):
    for cell in nb_data.get("cells", []):
        if cell.get("cell_type") == "code":
            yield "".join(cell.get("source", []))

def _iter_source_val_assigns(src):
    for node in _iter_source_assigns(src):
        try:
            val = python_util.ast_param_val(node.value)
        except TypeError:
            pass
        else:
            # Only support assigns to right-most name targets.
            target = node.targets[-1]
            if isinstance(target, ast.Name):
                yield node, target, node.value, val

def _iter_source_assigns(s):
    try:
        m = ast.parse(s)
    except Exception as e:
        if log.getEffectiveLevel() <= logging.DEBUG:
            log.exception("parsing %r", s)
        else:
            log.warning("error parsing Python source %r: %s", s, e)
    else:
        for node in m.body:
            if isinstance(node, ast.Assign):
                yield node

def _flag_data_for_val(val):
    return {
        "default": val,
        "type": _flag_type(val),
    }

def _flag_type(val):
    if isinstance(val, six.string_types):
        return "string"
    elif isinstance(val, bool):
        return "boolean"
    elif isinstance(val, (int, float)):
        return "number"
    else:
        return None
<|endoftext|>


<commit_before>import threading

import grpc

from six.moves import queue

import etcd3.etcdrpc as etcdrpc
import etcd3.events as events
import etcd3.exceptions as etcd3_exceptions
import etcd3.utils as utils

class Watch(object):

    def __init__(self, watch_id, iterator=None, etcd_client=None):
        self.watch_id = watch_id
        self.etcd_client = etcd_client
        self.iterator = iterator

    def cancel(self):
        self.etcd_client.cancel_watch(self.watch_id)

    def iterator(self):
        if self.iterator is not None:
            return self.iterator
        else:
            raise

class Watcher(threading.Thread):

    def __init__(self, watchstub, timeout=None, call_credentials=None, metadata=None):
        threading.Thread.__init__(self)
        self.timeout = timeout
        self._watch_id_callbacks = {}
        self._watch_id_queue = queue.Queue()
        self._watch_id_lock = threading.Lock()
        self._watch_requests_queue = queue.Queue()
        self._watch_response_iterator = watchstub.Watch(
            self._requests_iterator,
            credentials=call_credentials,
            metadata=metadata
        )
        self._callback = None
        self.daemon = True
        self.start()

    def run(self):
        try:
            for response in self._watch_response_iterator:
                if response.created:
                    self._watch_id_callbacks[response.watch_id] = \
                        self._callback
                    self._watch_id_queue.put(response.watch_id)

                callback = self._watch_id_callbacks.get(response.watch_id)
                if callback:
                    # The watcher can be safely reused, but adding a new event
                    # to indicate that the revision is already compacted
                    # requires api change which would break all users of this
                    # module. So, raising an exception if a watcher is still
                    # alive. The caller has to create a new client instance to
                    # recover would break all users of this module.
                    if response.compact_revision != 0:
                        callback(etcd3_exceptions.RevisionCompactedError(
                            response.compact_revision))
                        self.cancel(response.watch_id)
                        continue
                    for event in response.events:
                        callback(events.new_event(event))
        except grpc.RpcError as e:
            self.stop()
            if self._watch_id_callbacks:
                for callback in self._watch_id_callbacks.values():
                    callback(e)

    @property
    def _requests_iterator(self):
        while True:
            request, self._callback = self._watch_requests_queue.get()
            if request is None:
                break
            yield request

    def add_callback(self, key, callback,
                     range_end=None,
                     start_revision=None,
                     progress_notify=False,
                     filters=None,
                     prev_kv=False):
        with self._watch_id_lock:
            create_watch = etcdrpc.WatchCreateRequest()
            create_watch.key = utils.to_bytes(key)
            if range_end is not None:
                create_watch.range_end = utils.to_bytes(range_end)
            if start_revision is not None:
                create_watch.start_revision = start_revision
            if progress_notify:
                create_watch.progress_notify = progress_notify
            if filters is not None:
                create_watch.filters = filters
            if prev_kv:
                create_watch.prev_kv = prev_kv
            request = etcdrpc.WatchRequest(create_request=create_watch)
            self._watch_requests_queue.put((request, callback))
            return self._watch_id_queue.get(timeout=self.timeout)

    def cancel(self, watch_id):
        if watch_id is not None:
            self._watch_id_callbacks.pop(watch_id, None)
            cancel_watch = etcdrpc.WatchCancelRequest()
            cancel_watch.watch_id = watch_id
            request = etcdrpc.WatchRequest(cancel_request=cancel_watch)
            self._watch_requests_queue.put((request, None))

    def stop(self):
        self._watch_requests_queue.put((None, None))
<commit_msg>Lazy and fault tolerant watcher thread<commit_after>import logging
import threading

import attr
import grpc
import six
from six.moves import queue

import etcd3.etcdrpc as etcdrpc
import etcd3.events as events
import etcd3.exceptions as exceptions
import etcd3.utils as utils

_log = logging.getLogger(__name__)

class Watch(object):

    def __init__(self, watch_id, iterator=None, etcd_client=None):
        self.watch_id = watch_id
        self.etcd_client = etcd_client
        self.iterator = iterator

    def cancel(self):
        self.etcd_client.cancel_watch(self.watch_id)

    def iterator(self):
        if self.iterator is not None:
            return self.iterator

        raise ValueError('Undefined iterator')

class Watcher(object):

    def __init__(self, watchstub, timeout=None, call_credentials=None, metadata=None):
        self.timeout = timeout
        self._watch_stub = watchstub
        self._credentials = call_credentials
        self._metadata = metadata

        self._lock = threading.Lock()
        self._request_queue = queue.Queue(maxsize=10)
        self._callbacks = {}
        self._callback_thread = None
        self._new_watch_cond = threading.Condition(lock=self._lock)
        self._new_watch = None

    def add_callback(self, key, callback, range_end=None, start_revision=None,
                     progress_notify=False, filters=None, prev_kv=False):
        create_watch = etcdrpc.WatchCreateRequest()
        create_watch.key = utils.to_bytes(key)
        if range_end is not None:
            create_watch.range_end = utils.to_bytes(range_end)
        if start_revision is not None:
            create_watch.start_revision = start_revision
        if progress_notify:
            create_watch.progress_notify = progress_notify
        if filters is not None:
            create_watch.filters = filters
        if prev_kv:
            create_watch.prev_kv = prev_kv
        rq = etcdrpc.WatchRequest(create_request=create_watch)

        with self._lock:
            # Start the callback thread if it is not yet running.
            if not self._callback_thread:
                thread_name = 'etcd3_watch_%x' % (id(self),)
                self._callback_thread = threading.Thread(name=thread_name,
                                                         target=self._run)
                self._callback_thread.daemon = True
                self._callback_thread.start()

            # Only one create watch request can be pending at a time, so if
            # there one already, then wait for it to complete first.
            while self._new_watch:
                self._new_watch_cond.wait()

            # Submit a create watch request.
            new_watch = _NewWatch(callback)
            self._request_queue.put(rq)
            self._new_watch = new_watch

            # Wait for the request to be completed, or timeout.
            self._new_watch_cond.wait(timeout=self.timeout)
            self._new_watch = None

            # If the request not completed yet, then raise a timeout exception.
            if new_watch.id is None and new_watch.err is None:
                raise exceptions.WatchTimedOut()

            # Raise an exception if the watch request failed.
            if new_watch.err:
                raise new_watch.err

            # Wake up threads stuck on add_callback call if any.
            self._new_watch_cond.notify_all()
            return new_watch.id

    def cancel(self, watch_id):
        with self._lock:
            callback = self._callbacks.pop(watch_id, None)
            if not callback:
                return

            self._cancel_no_lock(watch_id)

    def _run(self):
        while True:
            response_iter = self._watch_stub.Watch(
                _new_request_iter(self._request_queue),
                credentials=self._credentials,
                metadata=self._metadata)
            try:
                for rs in response_iter:
                    self._handle_response(rs)

            except grpc.RpcError as err:
                with self._lock:
                    if self._new_watch:
                        self._new_watch.err = err
                        self._new_watch_cond.notify_all()

                    callbacks = self._callbacks
                    self._callbacks = {}

                    # Rotate request queue. This way we can terminate one gRPC
                    # stream and initiate another one whilst avoiding a race
                    # between them over requests in the queue.
                    self._request_queue.put(None)
                    self._request_queue = queue.Queue(maxsize=10)

                for callback in six.itervalues(callbacks):
                    _safe_callback(callback, err)

    def _handle_response(self, rs):
        with self._lock:
            if rs.created:
                # If the new watch request has already expired then cancel the
                # created watch right away.
                if not self._new_watch:
                    self._cancel_no_lock(rs.watch_id)
                    return

                if rs.compact_revision != 0:
                    self._new_watch.err = exceptions.RevisionCompactedError(
                        rs.compact_revision)
                    return

                self._callbacks[rs.watch_id] = self._new_watch.callback
                self._new_watch.id = rs.watch_id
                self._new_watch_cond.notify_all()

            callback = self._callbacks.get(rs.watch_id)

        # Ignore leftovers from canceled watches.
        if not callback:
            return

        # The watcher can be safely reused, but adding a new event
        # to indicate that the revision is already compacted
        # requires api change which would break all users of this
        # module. So, raising an exception if a watcher is still
        # alive.
        if rs.compact_revision != 0:
            err = exceptions.RevisionCompactedError(rs.compact_revision)
            _safe_callback(callback, err)
            self.cancel(rs.watch_id)
            return

        for event in rs.events:
            _safe_callback(callback, events.new_event(event))

    def _cancel_no_lock(self, watch_id):
        cancel_watch = etcdrpc.WatchCancelRequest()
        cancel_watch.watch_id = watch_id
        rq = etcdrpc.WatchRequest(cancel_request=cancel_watch)
        self._request_queue.put(rq)

@attr.s
class _NewWatch(object):
    callback = attr.ib()
    id = attr.ib(default=None)
    err = attr.ib(default=None)

def _new_request_iter(_request_queue):
    while True:
        rq = _request_queue.get()
        if rq is None:
            return

        yield rq

def _safe_callback(callback, event_or_err):
    try:
        callback(event_or_err)

    except Exception:
        _log.exception('Watch callback failed')
<|endoftext|>


<commit_before>from unittest import TestCase
from ass2m.cli import CLI
from tempfile import mkdtemp
import shutil
import sys
import re
from StringIO import StringIO

class BaseCLITest(TestCase):
    def setUp(self):
        self.root = mkdtemp(prefix='ass2m_test_root')
        self.app = CLI(self.root)

    def tearDown(self):
        if self.root:
            shutil.rmtree(self.root)

    def beginCapture(self):
        self.stdout = sys.stdout
        # begin capture
        sys.stdout = StringIO()

    def endCapture(self):
        captured = sys.stdout
        # end capture
        sys.stdout = self.stdout
        self.stdout = None
        return captured.getvalue()

    def test_init(self):
        self.beginCapture()
        assert self.app.main(['ass2m_test', 'init']) in (0, None)
        output = self.endCapture()
        assert output.strip() == "Ass2m working directory created."

        self.beginCapture()
        assert self.app.main(['ass2m_test', 'tree']) in (0, None)
        output = self.endCapture()
        assert re.match(re.escape(r'/')+r'\s+'+re.escape(r'all(rl-)'), output, re.S)
        assert re.match(".+"+re.escape(r'/.ass2m/')+r'\s+'+re.escape(r'all(---)'), output, re.S)
<commit_msg>Test for the automatic CLI finding of the root directory<commit_after>from unittest import TestCase
from ass2m.cli import CLI
from tempfile import mkdtemp
import os
import re
import shutil
import sys
from StringIO import StringIO

class BaseCLITest(TestCase):
    def setUp(self):
        self.root = mkdtemp(prefix='ass2m_test_root')
        self.app = CLI(self.root)

    def tearDown(self):
        if self.root:
            shutil.rmtree(self.root)

    def beginCapture(self):
        self.stdout = sys.stdout
        # begin capture
        sys.stdout = StringIO()

    def endCapture(self):
        captured = sys.stdout
        # end capture
        sys.stdout = self.stdout
        self.stdout = None
        return captured.getvalue()

    def test_init(self):
        self.beginCapture()
        assert self.app.main(['ass2m_test', 'init']) in (0, None)
        output = self.endCapture()
        assert output.strip() == "Ass2m working directory created."

        self.beginCapture()
        assert self.app.main(['ass2m_test', 'tree']) in (0, None)
        output = self.endCapture()
        assert re.match(re.escape(r'/')+r'\s+'+re.escape(r'all(rl-)'), output, re.S)
        assert re.match(".+"+re.escape(r'/.ass2m/')+r'\s+'+re.escape(r'all(---)'), output, re.S)

    def test_findRoot(self):
        self.beginCapture()
        assert self.app.main(['ass2m_test', 'init']) in (0, None)
        self.endCapture()

        subdir = os.path.join(self.root, "penguins")
        os.mkdir(subdir)
        subapp = CLI(subdir)
        assert subapp.main(['ass2m_test', 'tree']) in (0, None)
<|endoftext|>


<commit_before>import unittest

# noinspection PyClassHasNoInit
class Key:
    C = 0
    C_Sharp = 1
    D = 2
    D_Sharp = 3
    E = 4
    F = 5
    F_Sharp = 6
    G = 7
    G_Sharp = 8
    A = 9
    A_Sharp = 10
    B = 11

    @classmethod
    def all(cls):
        return range(12)

# Represents a note
class Note:
    def __init__(self, position, volume=112):
        self.position = position
        self.volume = volume

# Represents a chord
class Chord:
    # These lists are positions within a scale. [0, 2, 4] is first, third and fifth which is a regular triad
    triad = [0, 2, 4]
    triad_octave = [0, 2, 4, 7]
    suspended_second = [0, 2, 3]
    suspended_fourth = [0, 2, 5]
    seventh = [0, 2, 4, 7]
    seventh_octave = [0, 2, 4, 6, 7]
    sixth = [0, 2, 4, 5, 7]

    all = [triad, triad_octave, suspended_second, suspended_fourth, seventh, seventh_octave, sixth]

    def __init__(self, notes):
        self.notes = notes

    def add_stress(self, note):
        self.notes.append(note)

# Represents a scale
class Scale:
    major = [0, 2, 4, 5, 7, 9, 11]
    minor = [0, 2, 3, 5, 7, 8, 10]
    minor_pentatonic = [0, 3, 5, 7, 10]
    minor_blues = [0, 3, 5, 6, 7, 10]

    all = [major, minor, minor_pentatonic, minor_blues]

    # Make a new scale with a scale passed to it (e.g. scale = Scale(minor_blues))
    def __init__(self, scale, key=Key.C, base_octave=3):
        self.scale = scale
        self.length = len(scale)
        self.base_octave = base_octave
        self.key = key

    def interval_to_position(self, interval):
        return self.key + self.scale[interval % self.length] + 12 * (interval / self.length + self.base_octave)

    @classmethod
    def all_positions(cls, scale, key):
        s = Scale(scale, key, base_octave=0)
        interval = -7
        positions = []
        while True:
            pos = s.interval_to_position(interval)
            if pos > 127:
                break
            if pos >= 0:
                positions.append(pos)
            interval += 1
        return positions

    # Get a note from this scale starting with a given interval from the root (0, 1, 2, 3 etc.)
    def note(self, interval):
        position = self.interval_to_position(interval)
        return Note(position)

    # Get a chord from this scale starting with a given interval from the root (0, 1, 2, 3 etc.) Set the chord type
    # using intervals (e.g. chord = scale.chord(0, intervals=Chord.triad) gives the root triad. Chords always in key!)
    def chord(self, interval, intervals=Chord.triad):
        positions = map(lambda i: self.interval_to_position(interval + i), intervals)
        return Chord(map(Note, positions))

    # Go up by number of octaves
    def change_octave(self, by):
        self.base_octave = self.base_octave + by

keys_array = [set() for _ in range(128)]
for k in Key.all():
    for pos in Scale.all_positions(Scale.major, k):
        keys_array[pos].add(k)

def possible_keys(positions):
    key_set = set(range(12))
    for position in positions:
        key_set = key_set.intersection(keys_array[position])
    return key_set

class KeySelectionTestCase(unittest.TestCase):
    def test_keys_array(self):
        self.assertTrue(Key.C in keys_array[0])
        self.assertTrue(Key.C not in keys_array[1])

    def test_positions(self):
        self.assertTrue(0 in Scale.all_positions(Scale.major, Key.C))
        self.assertTrue(1 not in Scale.all_positions(Scale.major, Key.C))
        self.assertTrue(1 in Scale.all_positions(Scale.major, Key.C_Sharp))

    def test_possible_keys(self):
        self.assertTrue(Key.C in possible_keys([0, 2, 4]))
        self.assertTrue(Key.C_Sharp not in possible_keys([0, 2, 4]))

        self.assertTrue(Key.C not in possible_keys([1, 3, 5]))
        self.assertTrue(Key.C_Sharp in possible_keys([1, 3, 5]))

if __name__ == "__main__":
    unittest.main()
<commit_msg>key tracker queue determining kety<commit_after>import unittest
from Queue import Queue

# noinspection PyClassHasNoInit
class Key:
    C = 0
    C_Sharp = 1
    D = 2
    D_Sharp = 3
    E = 4
    F = 5
    F_Sharp = 6
    G = 7
    G_Sharp = 8
    A = 9
    A_Sharp = 10
    B = 11

    @classmethod
    def all(cls):
        return range(12)

# Represents a note
class Note:
    def __init__(self, position, volume=112):
        self.position = position
        self.volume = volume

# Represents a chord
class Chord:
    # These lists are positions within a scale. [0, 2, 4] is first, third and fifth which is a regular triad
    triad = [0, 2, 4]
    triad_octave = [0, 2, 4, 7]
    suspended_second = [0, 2, 3]
    suspended_fourth = [0, 2, 5]
    seventh = [0, 2, 4, 7]
    seventh_octave = [0, 2, 4, 6, 7]
    sixth = [0, 2, 4, 5, 7]

    all = [triad, triad_octave, suspended_second, suspended_fourth, seventh, seventh_octave, sixth]

    def __init__(self, notes):
        self.notes = notes

    def add_stress(self, note):
        self.notes.append(note)

# Represents a scale
class Scale:
    major = [0, 2, 4, 5, 7, 9, 11]
    minor = [0, 2, 3, 5, 7, 8, 10]
    minor_pentatonic = [0, 3, 5, 7, 10]
    minor_blues = [0, 3, 5, 6, 7, 10]

    all = [major, minor, minor_pentatonic, minor_blues]

    # Make a new scale with a scale passed to it (e.g. scale = Scale(minor_blues))
    def __init__(self, scale, key=Key.C, base_octave=3):
        self.scale = scale
        self.length = len(scale)
        self.base_octave = base_octave
        self.key = key

    def interval_to_position(self, interval):
        return self.key + self.scale[interval % self.length] + 12 * (interval / self.length + self.base_octave)

    @classmethod
    def all_positions(cls, scale, key):
        s = Scale(scale, key, base_octave=0)
        interval = -7
        positions = []
        while True:
            pos = s.interval_to_position(interval)
            if pos > 127:
                break
            if pos >= 0:
                positions.append(pos)
            interval += 1
        return positions

    # Get a note from this scale starting with a given interval from the root (0, 1, 2, 3 etc.)
    def note(self, interval):
        position = self.interval_to_position(interval)
        return Note(position)

    # Get a chord from this scale starting with a given interval from the root (0, 1, 2, 3 etc.) Set the chord type
    # using intervals (e.g. chord = scale.chord(0, intervals=Chord.triad) gives the root triad. Chords always in key!)
    def chord(self, interval, intervals=Chord.triad):
        positions = map(lambda i: self.interval_to_position(interval + i), intervals)
        return Chord(map(Note, positions))

    # Go up by number of octaves
    def change_octave(self, by):
        self.base_octave = self.base_octave + by

keys_array = [set() for _ in range(128)]
for k in Key.all():
    for pos in Scale.all_positions(Scale.major, k):
        keys_array[pos].add(k)

def possible_keys(positions):
    key_set = set(range(12))
    for position in positions:
        key_set = key_set.intersection(keys_array[position])
    return key_set

class KeyTracker:
    def __init__(self):
        self.queue = Queue()

    def add_note(self, note):
        self.queue.put(note)

    @property
    def keys(self):
        return possible_keys(list(self.queue.queue))

class KeySelectionTestCase(unittest.TestCase):
    def test_keys_array(self):
        self.assertTrue(Key.C in keys_array[0])
        self.assertTrue(Key.C not in keys_array[1])

    def test_positions(self):
        self.assertTrue(0 in Scale.all_positions(Scale.major, Key.C))
        self.assertTrue(1 not in Scale.all_positions(Scale.major, Key.C))
        self.assertTrue(1 in Scale.all_positions(Scale.major, Key.C_Sharp))

    def test_possible_keys(self):
        self.assertTrue(Key.C in possible_keys([0, 2, 4]))
        self.assertTrue(Key.C_Sharp not in possible_keys([0, 2, 4]))

        self.assertTrue(Key.C not in possible_keys([1, 3, 5]))
        self.assertTrue(Key.C_Sharp in possible_keys([1, 3, 5]))

    def test_key_tracker(self):
        key_tracker = KeyTracker()
        key_tracker.add_note(0)
        key_tracker.add_note(2)
        self.assertTrue(Key.C in key_tracker.keys)
        self.assertTrue(Key.G in key_tracker.keys)
        self.assertTrue(Key.C_Sharp not in key_tracker.keys)
        self.assertTrue(Key.D not in key_tracker.keys)

if __name__ == "__main__":
    unittest.main()
<|endoftext|>