<commit_before># -*- coding: UTF-8 -*-
from __future__ import absolute_import
from __future__ import unicode_literals

import io
import logging
import os.path
import shutil

import mock
import pytest

from pre_commit.staged_files_only import staged_files_only
from pre_commit.util import cmd_output
from pre_commit.util import cwd
from testing.auto_namedtuple import auto_namedtuple
from testing.fixtures import git_dir
from testing.util import get_resource_path

FOO_CONTENTS = '\n'.join(('1', '2', '3', '4', '5', '6', '7', '8', ''))

def get_short_git_status():
    git_status = cmd_output('git', 'status', '-s')[1]
    return dict(reversed(line.split()) for line in git_status.splitlines())

@pytest.yield_fixture
def foo_staged(tempdir_factory):
    path = git_dir(tempdir_factory)
    with cwd(path):
        with io.open('foo', 'w') as foo_file:
            foo_file.write(FOO_CONTENTS)
        cmd_output('git', 'add', 'foo')
        foo_filename = os.path.join(path, 'foo')
        yield auto_namedtuple(path=path, foo_filename=foo_filename)

def _test_foo_state(
        path,
        foo_contents=FOO_CONTENTS,
        status='A',
        encoding='UTF-8',
):
    assert os.path.exists(path.foo_filename)
    assert io.open(path.foo_filename, encoding=encoding).read() == foo_contents
    actual_status = get_short_git_status()['foo']
    assert status == actual_status

def test_foo_staged(foo_staged):
    _test_foo_state(foo_staged)

def test_foo_nothing_unstaged(foo_staged, cmd_runner):
    with staged_files_only(cmd_runner):
        _test_foo_state(foo_staged)
    _test_foo_state(foo_staged)

def test_foo_something_unstaged(foo_staged, cmd_runner):
    with io.open(foo_staged.foo_filename, 'w') as foo_file:
        foo_file.write('herp\nderp\n')

    _test_foo_state(foo_staged, 'herp\nderp\n', 'AM')

    with staged_files_only(cmd_runner):
        _test_foo_state(foo_staged)

    _test_foo_state(foo_staged, 'herp\nderp\n', 'AM')

def test_foo_something_unstaged_diff_color_always(foo_staged, cmd_runner):
    cmd_output('git', 'config', '--local', 'color.diff', 'always')
    test_foo_something_unstaged(foo_staged, cmd_runner)

def test_foo_both_modify_non_conflicting(foo_staged, cmd_runner):
    with io.open(foo_staged.foo_filename, 'w') as foo_file:
        foo_file.write(FOO_CONTENTS + '9\n')

    _test_foo_state(foo_staged, FOO_CONTENTS + '9\n', 'AM')

    with staged_files_only(cmd_runner):
        _test_foo_state(foo_staged)

        # Modify the file as part of the "pre-commit"
        with io.open(foo_staged.foo_filename, 'w') as foo_file:
            foo_file.write(FOO_CONTENTS.replace('1', 'a'))

        _test_foo_state(foo_staged, FOO_CONTENTS.replace('1', 'a'), 'AM')

    _test_foo_state(foo_staged, FOO_CONTENTS.replace('1', 'a') + '9\n', 'AM')

def test_foo_both_modify_conflicting(foo_staged, cmd_runner):
    with io.open(foo_staged.foo_filename, 'w') as foo_file:
        foo_file.write(FOO_CONTENTS.replace('1', 'a'))

    _test_foo_state(foo_staged, FOO_CONTENTS.replace('1', 'a'), 'AM')

    with staged_files_only(cmd_runner):
        _test_foo_state(foo_staged)

        # Modify in the same place as the stashed diff
        with io.open(foo_staged.foo_filename, 'w') as foo_file:
            foo_file.write(FOO_CONTENTS.replace('1', 'b'))

        _test_foo_state(foo_staged, FOO_CONTENTS.replace('1', 'b'), 'AM')

    _test_foo_state(foo_staged, FOO_CONTENTS.replace('1', 'a'), 'AM')

@pytest.yield_fixture
def img_staged(tempdir_factory):
    path = git_dir(tempdir_factory)
    with cwd(path):
        img_filename = os.path.join(path, 'img.jpg')
        shutil.copy(get_resource_path('img1.jpg'), img_filename)
        cmd_output('git', 'add', 'img.jpg')
        yield auto_namedtuple(path=path, img_filename=img_filename)

def _test_img_state(path, expected_file='img1.jpg', status='A'):
    assert os.path.exists(path.img_filename)
    assert (
        io.open(path.img_filename, 'rb').read() ==
        io.open(get_resource_path(expected_file), 'rb').read()
    )
    actual_status = get_short_git_status()['img.jpg']
    assert status == actual_status

def test_img_staged(img_staged):
    _test_img_state(img_staged)

def test_img_nothing_unstaged(img_staged, cmd_runner):
    with staged_files_only(cmd_runner):
        _test_img_state(img_staged)
    _test_img_state(img_staged)

def test_img_something_unstaged(img_staged, cmd_runner):
    shutil.copy(get_resource_path('img2.jpg'), img_staged.img_filename)

    _test_img_state(img_staged, 'img2.jpg', 'AM')

    with staged_files_only(cmd_runner):
        _test_img_state(img_staged)

    _test_img_state(img_staged, 'img2.jpg', 'AM')

def test_img_conflict(img_staged, cmd_runner):
    """Admittedly, this shouldn't happen, but just in case."""
    shutil.copy(get_resource_path('img2.jpg'), img_staged.img_filename)

    _test_img_state(img_staged, 'img2.jpg', 'AM')

    with staged_files_only(cmd_runner):
        _test_img_state(img_staged)
        shutil.copy(get_resource_path('img3.jpg'), img_staged.img_filename)
        _test_img_state(img_staged, 'img3.jpg', 'AM')

    _test_img_state(img_staged, 'img2.jpg', 'AM')

@pytest.yield_fixture
def submodule_with_commits(tempdir_factory):
    path = git_dir(tempdir_factory)
    with cwd(path):
        cmd_output('git', 'commit', '--allow-empty', '-m', 'foo')
        sha1 = cmd_output('git', 'rev-parse', 'HEAD')[1].strip()
        cmd_output('git', 'commit', '--allow-empty', '-m', 'bar')
        sha2 = cmd_output('git', 'rev-parse', 'HEAD')[1].strip()
        yield auto_namedtuple(path=path, sha1=sha1, sha2=sha2)

def checkout_submodule(sha):
    with cwd('sub'):
        cmd_output('git', 'checkout', sha)

@pytest.yield_fixture
def sub_staged(submodule_with_commits, tempdir_factory):
    path = git_dir(tempdir_factory)
    with cwd(path):
        cmd_output(
            'git', 'submodule', 'add', submodule_with_commits.path, 'sub',
        )
        checkout_submodule(submodule_with_commits.sha1)
        cmd_output('git', 'add', 'sub')
        yield auto_namedtuple(
            path=path,
            sub_path=os.path.join(path, 'sub'),
            submodule=submodule_with_commits,
        )

def _test_sub_state(path, sha='sha1', status='A'):
    assert os.path.exists(path.sub_path)
    with cwd(path.sub_path):
        actual_sha = cmd_output('git', 'rev-parse', 'HEAD')[1].strip()
    assert actual_sha == getattr(path.submodule, sha)
    actual_status = get_short_git_status()['sub']
    assert actual_status == status

def test_sub_staged(sub_staged):
    _test_sub_state(sub_staged)

def test_sub_nothing_unstaged(sub_staged, cmd_runner):
    with staged_files_only(cmd_runner):
        _test_sub_state(sub_staged)
    _test_sub_state(sub_staged)

def test_sub_something_unstaged(sub_staged, cmd_runner):
    checkout_submodule(sub_staged.submodule.sha2)

    _test_sub_state(sub_staged, 'sha2', 'AM')

    with staged_files_only(cmd_runner):
        # This is different from others, we don't want to touch subs
        _test_sub_state(sub_staged, 'sha2', 'AM')

    _test_sub_state(sub_staged, 'sha2', 'AM')

@pytest.yield_fixture
def fake_logging_handler():
    class FakeHandler(logging.Handler):
        def __init__(self):
            logging.Handler.__init__(self)
            self.logs = []

        def emit(self, record):
            self.logs.append(record)  # pragma: no cover (only hit in failure)

    pre_commit_logger = logging.getLogger('pre_commit')
    original_level = pre_commit_logger.getEffectiveLevel()
    handler = FakeHandler()
    pre_commit_logger.addHandler(handler)
    pre_commit_logger.setLevel(logging.WARNING)
    yield handler
    pre_commit_logger.setLevel(original_level)
    pre_commit_logger.removeHandler(handler)

def test_diff_returns_1_no_diff_though(fake_logging_handler, foo_staged):
    cmd_runner = mock.Mock()
    cmd_runner.run.return_value = (1, '', '')
    cmd_runner.path.return_value = '.pre-commit-files_patch'
    with staged_files_only(cmd_runner):
        pass
    assert not fake_logging_handler.logs

def test_stage_utf8_changes(foo_staged, cmd_runner):
    contents = '\u2603'
    with io.open('foo', 'w', encoding='UTF-8') as foo_file:
        foo_file.write(contents)

    _test_foo_state(foo_staged, contents, 'AM')
    with staged_files_only(cmd_runner):
        _test_foo_state(foo_staged)
    _test_foo_state(foo_staged, contents, 'AM')

def test_stage_non_utf8_changes(foo_staged, cmd_runner):
    contents = 'ú'
    # Produce a latin-1 diff
    with io.open('foo', 'w', encoding='latin-1') as foo_file:
        foo_file.write(contents)

    _test_foo_state(foo_staged, contents, 'AM', encoding='latin-1')
    with staged_files_only(cmd_runner):
        _test_foo_state(foo_staged)
    _test_foo_state(foo_staged, contents, 'AM', encoding='latin-1')

def test_non_utf8_conflicting_diff(foo_staged, cmd_runner):
    """Regression test for #397"""
    # The trailing whitespace is important here, this triggers git to produce
    # an error message which looks like:
    #
    # ...patch1471530032:14: trailing whitespace.
    # [[unprintable character]][[space character]]
    # error: patch failed: foo:1
    # error: foo: patch does not apply
    #
    # Previously, the error message (though discarded immediately) was being
    # decoded with the UTF-8 codec (causing a crash)
    contents = 'ú \n'
    with io.open('foo', 'w', encoding='latin-1') as foo_file:
        foo_file.write(contents)

    _test_foo_state(foo_staged, contents, 'AM', encoding='latin-1')
    with staged_files_only(cmd_runner):
        _test_foo_state(foo_staged)
        # Create a conflicting diff that will need to be rolled back
        with io.open('foo', 'w') as foo_file:
            foo_file.write('')
    _test_foo_state(foo_staged, contents, 'AM', encoding='latin-1')
<commit_msg>Add regression test for external diff tools<commit_after># -*- coding: UTF-8 -*-
from __future__ import absolute_import
from __future__ import unicode_literals

import io
import logging
import os.path
import shutil

import mock
import pytest

from pre_commit.staged_files_only import staged_files_only
from pre_commit.util import cmd_output
from pre_commit.util import cwd
from testing.auto_namedtuple import auto_namedtuple
from testing.fixtures import git_dir
from testing.util import get_resource_path

FOO_CONTENTS = '\n'.join(('1', '2', '3', '4', '5', '6', '7', '8', ''))

def get_short_git_status():
    git_status = cmd_output('git', 'status', '-s')[1]
    return dict(reversed(line.split()) for line in git_status.splitlines())

@pytest.yield_fixture
def foo_staged(tempdir_factory):
    path = git_dir(tempdir_factory)
    with cwd(path):
        with io.open('foo', 'w') as foo_file:
            foo_file.write(FOO_CONTENTS)
        cmd_output('git', 'add', 'foo')
        foo_filename = os.path.join(path, 'foo')
        yield auto_namedtuple(path=path, foo_filename=foo_filename)

def _test_foo_state(
        path,
        foo_contents=FOO_CONTENTS,
        status='A',
        encoding='UTF-8',
):
    assert os.path.exists(path.foo_filename)
    assert io.open(path.foo_filename, encoding=encoding).read() == foo_contents
    actual_status = get_short_git_status()['foo']
    assert status == actual_status

def test_foo_staged(foo_staged):
    _test_foo_state(foo_staged)

def test_foo_nothing_unstaged(foo_staged, cmd_runner):
    with staged_files_only(cmd_runner):
        _test_foo_state(foo_staged)
    _test_foo_state(foo_staged)

def test_foo_something_unstaged(foo_staged, cmd_runner):
    with io.open(foo_staged.foo_filename, 'w') as foo_file:
        foo_file.write('herp\nderp\n')

    _test_foo_state(foo_staged, 'herp\nderp\n', 'AM')

    with staged_files_only(cmd_runner):
        _test_foo_state(foo_staged)

    _test_foo_state(foo_staged, 'herp\nderp\n', 'AM')

def test_something_unstaged_ext_diff_tool(foo_staged, cmd_runner, tmpdir):
    diff_tool = tmpdir.join('diff-tool.sh')
    diff_tool.write('#!/usr/bin/env bash\necho "$@"\n')
    cmd_output('git', 'config', 'diff.external', diff_tool.strpath)
    test_foo_something_unstaged(foo_staged, cmd_runner)

def test_foo_something_unstaged_diff_color_always(foo_staged, cmd_runner):
    cmd_output('git', 'config', '--local', 'color.diff', 'always')
    test_foo_something_unstaged(foo_staged, cmd_runner)

def test_foo_both_modify_non_conflicting(foo_staged, cmd_runner):
    with io.open(foo_staged.foo_filename, 'w') as foo_file:
        foo_file.write(FOO_CONTENTS + '9\n')

    _test_foo_state(foo_staged, FOO_CONTENTS + '9\n', 'AM')

    with staged_files_only(cmd_runner):
        _test_foo_state(foo_staged)

        # Modify the file as part of the "pre-commit"
        with io.open(foo_staged.foo_filename, 'w') as foo_file:
            foo_file.write(FOO_CONTENTS.replace('1', 'a'))

        _test_foo_state(foo_staged, FOO_CONTENTS.replace('1', 'a'), 'AM')

    _test_foo_state(foo_staged, FOO_CONTENTS.replace('1', 'a') + '9\n', 'AM')

def test_foo_both_modify_conflicting(foo_staged, cmd_runner):
    with io.open(foo_staged.foo_filename, 'w') as foo_file:
        foo_file.write(FOO_CONTENTS.replace('1', 'a'))

    _test_foo_state(foo_staged, FOO_CONTENTS.replace('1', 'a'), 'AM')

    with staged_files_only(cmd_runner):
        _test_foo_state(foo_staged)

        # Modify in the same place as the stashed diff
        with io.open(foo_staged.foo_filename, 'w') as foo_file:
            foo_file.write(FOO_CONTENTS.replace('1', 'b'))

        _test_foo_state(foo_staged, FOO_CONTENTS.replace('1', 'b'), 'AM')

    _test_foo_state(foo_staged, FOO_CONTENTS.replace('1', 'a'), 'AM')

@pytest.yield_fixture
def img_staged(tempdir_factory):
    path = git_dir(tempdir_factory)
    with cwd(path):
        img_filename = os.path.join(path, 'img.jpg')
        shutil.copy(get_resource_path('img1.jpg'), img_filename)
        cmd_output('git', 'add', 'img.jpg')
        yield auto_namedtuple(path=path, img_filename=img_filename)

def _test_img_state(path, expected_file='img1.jpg', status='A'):
    assert os.path.exists(path.img_filename)
    assert (
        io.open(path.img_filename, 'rb').read() ==
        io.open(get_resource_path(expected_file), 'rb').read()
    )
    actual_status = get_short_git_status()['img.jpg']
    assert status == actual_status

def test_img_staged(img_staged):
    _test_img_state(img_staged)

def test_img_nothing_unstaged(img_staged, cmd_runner):
    with staged_files_only(cmd_runner):
        _test_img_state(img_staged)
    _test_img_state(img_staged)

def test_img_something_unstaged(img_staged, cmd_runner):
    shutil.copy(get_resource_path('img2.jpg'), img_staged.img_filename)

    _test_img_state(img_staged, 'img2.jpg', 'AM')

    with staged_files_only(cmd_runner):
        _test_img_state(img_staged)

    _test_img_state(img_staged, 'img2.jpg', 'AM')

def test_img_conflict(img_staged, cmd_runner):
    """Admittedly, this shouldn't happen, but just in case."""
    shutil.copy(get_resource_path('img2.jpg'), img_staged.img_filename)

    _test_img_state(img_staged, 'img2.jpg', 'AM')

    with staged_files_only(cmd_runner):
        _test_img_state(img_staged)
        shutil.copy(get_resource_path('img3.jpg'), img_staged.img_filename)
        _test_img_state(img_staged, 'img3.jpg', 'AM')

    _test_img_state(img_staged, 'img2.jpg', 'AM')

@pytest.yield_fixture
def submodule_with_commits(tempdir_factory):
    path = git_dir(tempdir_factory)
    with cwd(path):
        cmd_output('git', 'commit', '--allow-empty', '-m', 'foo')
        sha1 = cmd_output('git', 'rev-parse', 'HEAD')[1].strip()
        cmd_output('git', 'commit', '--allow-empty', '-m', 'bar')
        sha2 = cmd_output('git', 'rev-parse', 'HEAD')[1].strip()
        yield auto_namedtuple(path=path, sha1=sha1, sha2=sha2)

def checkout_submodule(sha):
    with cwd('sub'):
        cmd_output('git', 'checkout', sha)

@pytest.yield_fixture
def sub_staged(submodule_with_commits, tempdir_factory):
    path = git_dir(tempdir_factory)
    with cwd(path):
        cmd_output(
            'git', 'submodule', 'add', submodule_with_commits.path, 'sub',
        )
        checkout_submodule(submodule_with_commits.sha1)
        cmd_output('git', 'add', 'sub')
        yield auto_namedtuple(
            path=path,
            sub_path=os.path.join(path, 'sub'),
            submodule=submodule_with_commits,
        )

def _test_sub_state(path, sha='sha1', status='A'):
    assert os.path.exists(path.sub_path)
    with cwd(path.sub_path):
        actual_sha = cmd_output('git', 'rev-parse', 'HEAD')[1].strip()
    assert actual_sha == getattr(path.submodule, sha)
    actual_status = get_short_git_status()['sub']
    assert actual_status == status

def test_sub_staged(sub_staged):
    _test_sub_state(sub_staged)

def test_sub_nothing_unstaged(sub_staged, cmd_runner):
    with staged_files_only(cmd_runner):
        _test_sub_state(sub_staged)
    _test_sub_state(sub_staged)

def test_sub_something_unstaged(sub_staged, cmd_runner):
    checkout_submodule(sub_staged.submodule.sha2)

    _test_sub_state(sub_staged, 'sha2', 'AM')

    with staged_files_only(cmd_runner):
        # This is different from others, we don't want to touch subs
        _test_sub_state(sub_staged, 'sha2', 'AM')

    _test_sub_state(sub_staged, 'sha2', 'AM')

@pytest.yield_fixture
def fake_logging_handler():
    class FakeHandler(logging.Handler):
        def __init__(self):
            logging.Handler.__init__(self)
            self.logs = []

        def emit(self, record):
            self.logs.append(record)  # pragma: no cover (only hit in failure)

    pre_commit_logger = logging.getLogger('pre_commit')
    original_level = pre_commit_logger.getEffectiveLevel()
    handler = FakeHandler()
    pre_commit_logger.addHandler(handler)
    pre_commit_logger.setLevel(logging.WARNING)
    yield handler
    pre_commit_logger.setLevel(original_level)
    pre_commit_logger.removeHandler(handler)

def test_diff_returns_1_no_diff_though(fake_logging_handler, foo_staged):
    cmd_runner = mock.Mock()
    cmd_runner.run.return_value = (1, '', '')
    cmd_runner.path.return_value = '.pre-commit-files_patch'
    with staged_files_only(cmd_runner):
        pass
    assert not fake_logging_handler.logs

def test_stage_utf8_changes(foo_staged, cmd_runner):
    contents = '\u2603'
    with io.open('foo', 'w', encoding='UTF-8') as foo_file:
        foo_file.write(contents)

    _test_foo_state(foo_staged, contents, 'AM')
    with staged_files_only(cmd_runner):
        _test_foo_state(foo_staged)
    _test_foo_state(foo_staged, contents, 'AM')

def test_stage_non_utf8_changes(foo_staged, cmd_runner):
    contents = 'ú'
    # Produce a latin-1 diff
    with io.open('foo', 'w', encoding='latin-1') as foo_file:
        foo_file.write(contents)

    _test_foo_state(foo_staged, contents, 'AM', encoding='latin-1')
    with staged_files_only(cmd_runner):
        _test_foo_state(foo_staged)
    _test_foo_state(foo_staged, contents, 'AM', encoding='latin-1')

def test_non_utf8_conflicting_diff(foo_staged, cmd_runner):
    """Regression test for #397"""
    # The trailing whitespace is important here, this triggers git to produce
    # an error message which looks like:
    #
    # ...patch1471530032:14: trailing whitespace.
    # [[unprintable character]][[space character]]
    # error: patch failed: foo:1
    # error: foo: patch does not apply
    #
    # Previously, the error message (though discarded immediately) was being
    # decoded with the UTF-8 codec (causing a crash)
    contents = 'ú \n'
    with io.open('foo', 'w', encoding='latin-1') as foo_file:
        foo_file.write(contents)

    _test_foo_state(foo_staged, contents, 'AM', encoding='latin-1')
    with staged_files_only(cmd_runner):
        _test_foo_state(foo_staged)
        # Create a conflicting diff that will need to be rolled back
        with io.open('foo', 'w') as foo_file:
            foo_file.write('')
    _test_foo_state(foo_staged, contents, 'AM', encoding='latin-1')
<|endoftext|>


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

"""Simple interface between Slack Webhooks and Zotero API

   Intended to be used in a cron-like system that will periodically
   run this script with since=`last run's version` which can be
   obtained using the artifact parameter's functionality

   Example: slack-zotero.py --group 12345 \
                            --api "kjas734890fnlkjafFJnadsf" \
                            --webhook "https://hooks.slack.com/services/OUHAEFNF/OUIHAQEUN/uihsdf786SHBF6ebSF" \
                            --since 9043 \
                            --channel "#test" \
                            --username "Zotero Bot" \
                            --icon ":cow:" \
                            --artifact "slack-zotero-bot-previous.json" \
                            -v
"""

import urllib.request
import json
import requests
import re
import html
import time
import datetime
import sys

def retrieve_articles(group_id, api_key, limit=1, include='data', since=0):
    """Retrieves articles from a Zotero group API feed"""
    zotero_template = "https://api.zotero.org/groups/{group_id}/items/top?start=0&limit={limit}&format=json&v=3&key={api_key}"
    zotero_url = zotero_template.format(group_id=group_id, api_key=api_key, limit=limit)

    zotero_url += "&include={0}".format(include) if include else ""
    zotero_url += "&since={0}".format(since) if since else ""

    print("Retrieving most recent {limit} articles since version: {version}".format(limit=limit, version=since))

    response = urllib.request.urlopen(zotero_url)
    articles = json.loads(response.readall().decode('utf-8'))

    print("Retrieved {0} articles".format(len(articles)))

    return articles

def send_article_to_slack(webhook_url, article, channel=None, username=None, icon_emoji=None, verbose=True, mock=False):
    """Sends a JSON article to the given Slack Webhooks URL"""
    payload = {"text": format_article(article)}
    if channel:
        payload['channel'] = channel

    if username:
        payload['username'] = username

    if icon_emoji:
        payload['icon_emoji'] = icon_emoji

    response = None
    if not mock:
        response = requests.post(webhook_url, json=payload)
        if response.status_code != 200:
            print(
                'Request to slack returned an error %s, the response is:\n%s'
                % (response.status_code, response.text)
            )

    if verbose:
        version = int(article['version'])
        print("{0} - {1}".format(version, article['data']['title']))

    return response

def format_article(article):
    """Feel free to overwrite me with your preferred format"""
    data = article['data']
    meta = article['meta']

    title = data['title']

    submitter = meta.get('createdByUser', {}).get('username', '')

    itemType = data.get('itemType', '')
    if itemType == "thesis":
        journal = data.get('university', '')
    else:
        journal = data.get('publicationTitle', '')

    authors = meta.get('creatorSummary', '').rstrip(".")
    date = data.get('date', '')
    if date or journal or authors:
        citation = "{authors}. _{journal}_ {date}".format(date=date, journal=journal, authors=authors)
    else:
        citation = ""

    abstract = data.get('abstractNote', '')
    if abstract:
        # Extract first N words
        word_cnt = 100
        abstract_words = abstract.split(" ")
        abstract = " ".join(abstract_words[:word_cnt])
        if len(abstract_words) > word_cnt:
            abstract += " …"

    url = data.get('url', '').strip()
    doi = data.get('DOI', '')
    tags = [t['tag'] for t in data.get('tags', '')]

    link = ""
    if not doi and url:
        link = url
    elif doi:
        link = "https://doi.org/{0}".format(doi)

    template = ""
    template += "<{link}|*{title}*>\n" if link else "*{title}*\n"
    template += "*Citation:* {citation}\n" if citation else ""
    template += "*Tags:* {tags}\n" if tags else ""
    template += "*Added By:* {submitter}\n" if submitter else ""

    template += "\n*Abstract:*\n```{abstract}```" if abstract else ""

    return template.format(title=title, abstract=abstract, link=link, submitter=submitter,
                           citation=citation, tags=", ".join(tags))

def main(zotero_group, zotero_api_key, slack_webhook_url, since_version=0, channel=None, username=None, icon_emoji=None,
         limit=25, mock=False, verbose=True):
    timestamp = datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d %H:%M:%S')

    articles = retrieve_articles(zotero_group, zotero_api_key, limit=limit if since_version else 1, since=since_version)

    max_version = max([since_version] + [article['version'] for article in articles])

    skipped = 0
    for article in reversed(articles):
        try:
            send_article_to_slack(slack_webhook_url, article, channel=channel, username=username, icon_emoji=icon_emoji,
                                  verbose=verbose, mock=mock)
        except Exception as inst:
            skipped += 1
            print("Problem Sending article {key}: {error}".format(key=article['key'], error=inst))

    print("Current Version: {0}".format(max_version))

    run_info = {"time": timestamp, "version": max_version, "articles_cnt": len(articles), "skipped": skipped}

    return run_info

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description="Retrieve most recent articles from a Zotero group using Zotero's "
                                                 "API v3 and send them to a Slack channel using Webhooks")

    parser.add_argument('--group', dest='group', type=int, required=False,
                        help='Zotero group ID of the library to monitor')

    parser.add_argument('--api', dest='api_key', type=str, required=False,
                        help='Zotero API key with access to the library')

    parser.add_argument('--webhook', dest='webhook', type=str, required=False,
                        help='Slack webhook URL to send articles')

    parser.add_argument('--since', dest='version', type=int, required=False, default=0,
                        help='Retrieve only articles created after this version')
    parser.add_argument('--limit', dest='limit', type=int, required=False, default=25,
                        help='Max articles to return when --since is supplied')
    parser.add_argument('--channel', dest='channel', type=str, required=False, default=None,
                        help='Override default Slack webhooks channel')
    parser.add_argument('--username', dest='username', type=str, required=False, default=None,
                        help='Override default Slack webhooks username')
    parser.add_argument('--icon', dest='icon_emoji', type=str, required=False, default=None,
                        help='Override default Slack webhooks icon')

    parser.add_argument('--artifact', dest='artifact', type=str, required=False, default=None,
                        help='Retrieve --since from & write run info to this file. OVERRIDES --since')

    parser.add_argument('--mock', dest='mock', action='store_true',
                        help='Mock run; will not write to Slack or filesystem')

    parser.add_argument('-v', dest='verbose', action='store_true',
                        help='More verbose logging')

    parser.add_argument('--test', dest='test', type=str, required=False, default=None,
                        help='Test output using JSON file, will not write to Slack or filesystem')

    args = parser.parse_args()

    since = args.version
    if args.artifact:
        try:
            with open(args.artifact) as data_file:
                data = json.load(data_file)
                since = data['version']
        except (FileNotFoundError, KeyError) as e:
            print("Error reading version info from artifact file, defaulting to {0}.".format(since))

    if args.test:
        # Monkey patching

        # Just in case
        args.mock = True

        try:
            with open(args.test) as data_file:
                articles = json.load(data_file)
        except (FileNotFoundError, KeyError) as e:
            print("Error reading test file.")
            articles = []

        def mocked_retrieve(*args, **kwargs):
            return articles

        retrieve_articles = mocked_retrieve

        def mocked_send(*args, **kwargs):
            print(format_article(args[1]))
            print("-------------------")

        send_article_to_slack = mocked_send

    run_info = main(args.group, args.api_key, args.webhook, since, args.channel, args.username, args.icon_emoji,
                    args.limit, args.mock, args.verbose)

    if not args.mock and not args.test and args.artifact:
        print("Writing run version to {artifact}".format(artifact=args.artifact))
        with open(args.artifact, 'w') as outfile:
            json.dump(run_info, outfile)

    if run_info['skipped']:
        # Articles were skipped, warn cron-like process by exiting with 2
        sys.exit(2)
<commit_msg>Format: Improve citation builder<commit_after>#!/usr/bin/env python

"""Simple interface between Slack Webhooks and Zotero API

   Intended to be used in a cron-like system that will periodically
   run this script with since=`last run's version` which can be
   obtained using the artifact parameter's functionality

   Example: slack-zotero.py --group 12345 \
                            --api "kjas734890fnlkjafFJnadsf" \
                            --webhook "https://hooks.slack.com/services/OUHAEFNF/OUIHAQEUN/uihsdf786SHBF6ebSF" \
                            --since 9043 \
                            --channel "#test" \
                            --username "Zotero Bot" \
                            --icon ":cow:" \
                            --artifact "slack-zotero-bot-previous.json" \
                            -v
"""

import urllib.request
import json
import requests
import re
import html
import time
import datetime
import sys

def retrieve_articles(group_id, api_key, limit=1, include='data', since=0):
    """Retrieves articles from a Zotero group API feed"""
    zotero_template = "https://api.zotero.org/groups/{group_id}/items/top?start=0&limit={limit}&format=json&v=3&key={api_key}"
    zotero_url = zotero_template.format(group_id=group_id, api_key=api_key, limit=limit)

    zotero_url += "&include={0}".format(include) if include else ""
    zotero_url += "&since={0}".format(since) if since else ""

    print("Retrieving most recent {limit} articles since version: {version}".format(limit=limit, version=since))

    response = urllib.request.urlopen(zotero_url)
    articles = json.loads(response.readall().decode('utf-8'))

    print("Retrieved {0} articles".format(len(articles)))

    return articles

def send_article_to_slack(webhook_url, article, channel=None, username=None, icon_emoji=None, verbose=True, mock=False):
    """Sends a JSON article to the given Slack Webhooks URL"""
    payload = {"text": format_article(article)}
    if channel:
        payload['channel'] = channel

    if username:
        payload['username'] = username

    if icon_emoji:
        payload['icon_emoji'] = icon_emoji

    response = None
    if not mock:
        response = requests.post(webhook_url, json=payload)
        if response.status_code != 200:
            print(
                'Request to slack returned an error %s, the response is:\n%s'
                % (response.status_code, response.text)
            )

    if verbose:
        version = int(article['version'])
        print("{0} - {1}".format(version, article['data']['title']))

    return response

def format_article(article):
    """Feel free to overwrite me with your preferred format"""
    data = article['data']
    meta = article['meta']

    title = data['title']

    submitter = meta.get('createdByUser', {}).get('username', '')

    itemType = data.get('itemType', '')
    if itemType == "thesis":
        journal = data.get('university', '')
    else:
        journal = data.get('publicationTitle', '')

    authors = meta.get('creatorSummary', '').rstrip(".")
    date = data.get('date', '')

    # Build citation
    citation = ""
    citation += (authors + ". ") if authors else ""
    citation += ("_" + journal + "_ ") if journal else ""
    citation += date if date else ""

    abstract = data.get('abstractNote', '')
    if abstract:
        # Extract first N words
        word_cnt = 100
        abstract_words = abstract.split(" ")
        abstract = " ".join(abstract_words[:word_cnt])
        if len(abstract_words) > word_cnt:
            abstract += " …"

    url = data.get('url', '').strip()
    doi = data.get('DOI', '')
    tags = [t['tag'] for t in data.get('tags', '')]

    link = ""
    if not doi and url:
        link = url
    elif doi:
        link = "https://doi.org/{0}".format(doi)

    template = ""
    template += "<{link}|*{title}*>\n" if link else "*{title}*\n"
    template += "*Citation:* {citation}\n" if citation else ""
    template += "*Tags:* {tags}\n" if tags else ""
    template += "*Added By:* {submitter}\n" if submitter else ""

    template += "\n*Abstract:*\n```{abstract}```" if abstract else ""

    return template.format(title=title, abstract=abstract, link=link, submitter=submitter,
                           citation=citation, tags=", ".join(tags))

def main(zotero_group, zotero_api_key, slack_webhook_url, since_version=0, channel=None, username=None, icon_emoji=None,
         limit=25, mock=False, verbose=True):
    timestamp = datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d %H:%M:%S')

    articles = retrieve_articles(zotero_group, zotero_api_key, limit=limit if since_version else 1, since=since_version)

    max_version = max([since_version] + [article['version'] for article in articles])

    skipped = 0
    for article in reversed(articles):
        try:
            send_article_to_slack(slack_webhook_url, article, channel=channel, username=username, icon_emoji=icon_emoji,
                                  verbose=verbose, mock=mock)
        except Exception as inst:
            skipped += 1
            print("Problem Sending article {key}: {error}".format(key=article['key'], error=inst))

    print("Current Version: {0}".format(max_version))

    run_info = {"time": timestamp, "version": max_version, "articles_cnt": len(articles), "skipped": skipped}

    return run_info

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description="Retrieve most recent articles from a Zotero group using Zotero's "
                                                 "API v3 and send them to a Slack channel using Webhooks")

    parser.add_argument('--group', dest='group', type=int, required=False,
                        help='Zotero group ID of the library to monitor')

    parser.add_argument('--api', dest='api_key', type=str, required=False,
                        help='Zotero API key with access to the library')

    parser.add_argument('--webhook', dest='webhook', type=str, required=False,
                        help='Slack webhook URL to send articles')

    parser.add_argument('--since', dest='version', type=int, required=False, default=0,
                        help='Retrieve only articles created after this version')
    parser.add_argument('--limit', dest='limit', type=int, required=False, default=25,
                        help='Max articles to return when --since is supplied')
    parser.add_argument('--channel', dest='channel', type=str, required=False, default=None,
                        help='Override default Slack webhooks channel')
    parser.add_argument('--username', dest='username', type=str, required=False, default=None,
                        help='Override default Slack webhooks username')
    parser.add_argument('--icon', dest='icon_emoji', type=str, required=False, default=None,
                        help='Override default Slack webhooks icon')

    parser.add_argument('--artifact', dest='artifact', type=str, required=False, default=None,
                        help='Retrieve --since from & write run info to this file. OVERRIDES --since')

    parser.add_argument('--mock', dest='mock', action='store_true',
                        help='Mock run; will not write to Slack or filesystem')

    parser.add_argument('-v', dest='verbose', action='store_true',
                        help='More verbose logging')

    parser.add_argument('--test', dest='test', type=str, required=False, default=None,
                        help='Test output using JSON file, will not write to Slack or filesystem')

    args = parser.parse_args()

    since = args.version
    if args.artifact:
        try:
            with open(args.artifact) as data_file:
                data = json.load(data_file)
                since = data['version']
        except (FileNotFoundError, KeyError) as e:
            print("Error reading version info from artifact file, defaulting to {0}.".format(since))

    if args.test:
        # Monkey patching

        # Just in case
        args.mock = True

        try:
            with open(args.test) as data_file:
                articles = json.load(data_file)
        except (FileNotFoundError, KeyError) as e:
            print("Error reading test file.")
            articles = []

        def mocked_retrieve(*args, **kwargs):
            return articles

        retrieve_articles = mocked_retrieve

        def mocked_send(*args, **kwargs):
            print(format_article(args[1]))
            print("-------------------")

        send_article_to_slack = mocked_send

    run_info = main(args.group, args.api_key, args.webhook, since, args.channel, args.username, args.icon_emoji,
                    args.limit, args.mock, args.verbose)

    if not args.mock and not args.test and args.artifact:
        print("Writing run version to {artifact}".format(artifact=args.artifact))
        with open(args.artifact, 'w') as outfile:
            json.dump(run_info, outfile)

    if run_info['skipped']:
        # Articles were skipped, warn cron-like process by exiting with 2
        sys.exit(2)
<|endoftext|>


<commit_before>INSTALLED_APPS = (
    'nexmo',
)
DATABASES = {
    'default': {
        'ENGINE': 'django.db.backends.sqlite3',
        'NAME': ':memory:',
    }
}
MIDDLEWARE_CLASSES = (
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.common.CommonMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
    'django.middleware.clickjacking.XFrameOptionsMiddleware',
)
SECRET_KEY = "gloubiboulga secret key"
<commit_msg>Dummy nexmo test settings<commit_after>INSTALLED_APPS = (
    'nexmo',
)
DATABASES = {
    'default': {
        'ENGINE': 'django.db.backends.sqlite3',
        'NAME': ':memory:',
    }
}
MIDDLEWARE_CLASSES = (
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.common.CommonMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
    'django.middleware.clickjacking.XFrameOptionsMiddleware',
)
SECRET_KEY = "gloubiboulga secret key"
TEST_RUNNER = 'django.test.runner.DiscoverRunner'

NEXMO_API_KEY = 'API_KEY'
NEXMO_API_SECRET = 'SECRET'
NEXMO_DEFAULT_FROM = '+33123456789'
<|endoftext|>


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

# project moved to: https://github.com/itzo/slack-beebot

import websocket
import os
import time
from slackclient import SlackClient
import sqlite3 as db
import sys
import re

channel = '#general'
users = {}

# add timestamp to all print statements
old_out = sys.stdout
class timestamped:
    nl = True

    def write(self, x):
        # overload write()
        if x == '\n':
            old_out.write(x)
            self.nl = True
        elif self.nl:
            old_out.write('[ %s ] %s' % (str(int(time.time())), x))
            self.nl = False
        else:
            old_out.write(x)
sys.stdout = timestamped()

# crate db table if none exists
def create_db():
    try:
        con = db.connect('reactions.db')
        cur = con.cursor()
        cur.execute('DROP TABLE IF EXISTS reactions')
        cur.executescript("""
            CREATE TABLE reactions(
                from_user TEXT,
                to_user TEXT,
                reaction TEXT,
                counter INTEGER
            );""")
        con.commit()
    except db.Error, e:
        if con:
            con.rollback()
        print "Error %s:" % e.args[0]
        sys.exit(1)
    finally:
        if con:
            con.close()

# insert reaction into db table
def db_insert(from_user, to_user, reaction, counter):
    if os.path.isfile('reactions.db'):
        con = db.connect('reactions.db')
        cur = con.cursor()
        cur.execute('INSERT INTO reactions VALUES(?,?,?,?)', \
            (from_user, to_user, reaction, counter));
        con.commit()
    else:
        print "Can't find the database.\n"
        sys.exit(2)

# parse slack events for reactions / commands
def parse_event(event):
    #print str(event) + '\n'
    if event and len(event) > 0:

        # process reactions
        data = event[0]
        if all(x in data for x in ['user', 'item_user', 'reaction']):
            reaction = data['reaction']
            from_user = data['user']
            to_user = data['item_user']
            if data['type'] == 'reaction_added' and from_user != to_user:
                print "%s reacted with '%s' to %s" % (users[from_user], reaction, users[to_user])
                counter = '1'
                db_insert(from_user, to_user, reaction, counter)
            elif data['type'] == 'reaction_removed' and from_user != to_user:
                print "%s withdrew their reaction of '%s' from %s" % (users[from_user], reaction, users[to_user])
                counter = '-1'
                db_insert(from_user, to_user, reaction, counter)
            return reaction, from_user, to_user

        # answer commands
        if 'text' in data:
            if data['text'].lower().startswith('showme top'):
                channel_id = data['channel']
                if len(data['text'].split()) > 2:
                    reaction = data['text'].lower().split()[2]
                    if re.match(r'^[A-Za-z0-9_+]+$', reaction):
                        print_top(reaction, channel_id)
                        #if len(response) < 1:
                        #    response = 'Not sure we have these stats yet...'
                        #sc.api_call("chat.postMessage", channel=channel, text=response, as_user=True)
                    else:
                        bot_usage(channel_id)
                else:
                    bot_usage(channel_id)

    return None, None, None

# send a message with the correct way to use the bot
def bot_usage(channel_id):
    sc.api_call("chat.postMessage", channel=channel_id, text='usage: showme top <reaction>', as_user=True)

# print top recipients of a reaction
def print_top(reaction, channel_id):
    if os.path.isfile('reactions.db'):
        con = db.connect('reactions.db')
        with con:
            cur = con.cursor()
            cur.execute("SELECT to_user, sum(counter) as count from reactions where reaction='"+ reaction +"' group by to_user order by count desc limit 5;")
            con.commit()
            rows = cur.fetchall()
            print "Showing top "+reaction
            response = "```"
            if len(rows) > 0:
                for row in rows:
                    print "%-14s %+14d" % (users[row[0]], row[1])
                    response += "{:14} {:14d}\n".format(users[row[0]],row[1])
            else:
                response += "no '"+reaction+"' reactions found"
                print "none found"
            response += "```"
            sc.api_call("chat.postMessage", channel=channel_id, text=response, as_user=True)
    else:
        print "Can't find the database.\n"
        sys.exit(2)

# get the list of users and their names for later use
def get_users():
    data = sc.api_call('users.list', channel='#general')
    for user in data['members']:
        print 'id: %s, name: %s' % (user['id'], user['name'])
        users[user['id']] = user['name']

# main
if __name__ == '__main__':
    # initialize db if it doesn't exist
    if os.path.exists('./reactions.db') == False:
        create_db()
    # connect to slack
    token = os.environ.get('SLACK_BOT_TOKEN')
    sc = SlackClient(token)
    try:
        if sc.rtm_connect():
            print('Bot connected and running!')
            #sc.api_call("chat.postMessage", channel=channel, text="beebot is back from the dead...", as_user=True)
            get_users()
            while True:
                reaction, from_user, to_user = parse_event(sc.rtm_read())
                time.sleep(1)
        else:
            print 'Connection failed. Check token.'
    except websocket._exceptions.WebSocketConnectionClosedException:
        print 'Connection closed. Did someone disable the bot integration?'

<commit_msg>Moved to: https://github.com/itzo/slack-beebot<commit_after>#!/usr/bin/env python

#
# Project moved to: https://github.com/itzo/slack-beebot
# This is just kept for historical reasons
#

import websocket
import os
import time
from slackclient import SlackClient
import sqlite3 as db
import sys
import re

channel = '#general'
users = {}

# add timestamp to all print statements
old_out = sys.stdout
class timestamped:
    nl = True

    def write(self, x):
        # overload write()
        if x == '\n':
            old_out.write(x)
            self.nl = True
        elif self.nl:
            old_out.write('[ %s ] %s' % (str(int(time.time())), x))
            self.nl = False
        else:
            old_out.write(x)
sys.stdout = timestamped()

# crate db table if none exists
def create_db():
    try:
        con = db.connect('reactions.db')
        cur = con.cursor()
        cur.execute('DROP TABLE IF EXISTS reactions')
        cur.executescript("""
            CREATE TABLE reactions(
                from_user TEXT,
                to_user TEXT,
                reaction TEXT,
                counter INTEGER
            );""")
        con.commit()
    except db.Error, e:
        if con:
            con.rollback()
        print "Error %s:" % e.args[0]
        sys.exit(1)
    finally:
        if con:
            con.close()

# insert reaction into db table
def db_insert(from_user, to_user, reaction, counter):
    if os.path.isfile('reactions.db'):
        con = db.connect('reactions.db')
        cur = con.cursor()
        cur.execute('INSERT INTO reactions VALUES(?,?,?,?)', \
            (from_user, to_user, reaction, counter));
        con.commit()
    else:
        print "Can't find the database.\n"
        sys.exit(2)

# parse slack events for reactions / commands
def parse_event(event):
    #print str(event) + '\n'
    if event and len(event) > 0:

        # process reactions
        data = event[0]
        if all(x in data for x in ['user', 'item_user', 'reaction']):
            reaction = data['reaction']
            from_user = data['user']
            to_user = data['item_user']
            if data['type'] == 'reaction_added' and from_user != to_user:
                print "%s reacted with '%s' to %s" % (users[from_user], reaction, users[to_user])
                counter = '1'
                db_insert(from_user, to_user, reaction, counter)
            elif data['type'] == 'reaction_removed' and from_user != to_user:
                print "%s withdrew their reaction of '%s' from %s" % (users[from_user], reaction, users[to_user])
                counter = '-1'
                db_insert(from_user, to_user, reaction, counter)
            return reaction, from_user, to_user

        # answer commands
        if 'text' in data:
            if data['text'].lower().startswith('showme top'):
                channel_id = data['channel']
                if len(data['text'].split()) > 2:
                    reaction = data['text'].lower().split()[2]
                    if re.match(r'^[A-Za-z0-9_+]+$', reaction):
                        print_top(reaction, channel_id)
                        #if len(response) < 1:
                        #    response = 'Not sure we have these stats yet...'
                        #sc.api_call("chat.postMessage", channel=channel, text=response, as_user=True)
                    else:
                        bot_usage(channel_id)
                else:
                    bot_usage(channel_id)

    return None, None, None

# send a message with the correct way to use the bot
def bot_usage(channel_id):
    sc.api_call("chat.postMessage", channel=channel_id, text='usage: showme top <reaction>', as_user=True)

# print top recipients of a reaction
def print_top(reaction, channel_id):
    if os.path.isfile('reactions.db'):
        con = db.connect('reactions.db')
        with con:
            cur = con.cursor()
            cur.execute("SELECT to_user, sum(counter) as count from reactions where reaction='"+ reaction +"' group by to_user order by count desc limit 5;")
            con.commit()
            rows = cur.fetchall()
            print "Showing top "+reaction
            response = "```"
            if len(rows) > 0:
                for row in rows:
                    print "%-14s %+14d" % (users[row[0]], row[1])
                    response += "{:14} {:14d}\n".format(users[row[0]],row[1])
            else:
                response += "no '"+reaction+"' reactions found"
                print "none found"
            response += "```"
            sc.api_call("chat.postMessage", channel=channel_id, text=response, as_user=True)
    else:
        print "Can't find the database.\n"
        sys.exit(2)

# get the list of users and their names for later use
def get_users():
    data = sc.api_call('users.list', channel='#general')
    for user in data['members']:
        print 'id: %s, name: %s' % (user['id'], user['name'])
        users[user['id']] = user['name']

# main
if __name__ == '__main__':
    # initialize db if it doesn't exist
    if os.path.exists('./reactions.db') == False:
        create_db()
    # connect to slack
    token = os.environ.get('SLACK_BOT_TOKEN')
    sc = SlackClient(token)
    try:
        if sc.rtm_connect():
            print('Bot connected and running!')
            #sc.api_call("chat.postMessage", channel=channel, text="beebot is back from the dead...", as_user=True)
            get_users()
            while True:
                reaction, from_user, to_user = parse_event(sc.rtm_read())
                time.sleep(1)
        else:
            print 'Connection failed. Check token.'
    except websocket._exceptions.WebSocketConnectionClosedException:
        print 'Connection closed. Did someone disable the bot integration?'

<|endoftext|>


<commit_before># Copyright (c) 2016 Philip Hane
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

from .exceptions import (IPTablesError, IPTablesExists, IPTablesNotExists,
                         SubprocessError)
from .utils import popen_wrapper
import logging

log = logging.getLogger(__name__)

class IPTablesSinkhole:
    """
    The class for managing sinkhole configuration within iptables.

    Args:
        interface: The secondary network interface dedicated to sinkhole
            traffic.
            Warning: Do not accidentally set this to your primary interface.
            It will drop all traffic, and kill your remote access.
        interface_addr: The IP address assigned to interface.
        log_prefix: Prefix for syslog messages.
        protocol: The protocol(s) to log (all traffic will still be dropped).
            Accepts a comma separated string of protocols
            (tcp,udp,udplite,icmp,esp,ah,sctp) or all.
        dport: The destination port(s) to log (for applicable protocols).
            Range should be in the format startport:endport or 0,1,2,3,n..
        hashlimit: Set the hashlimit rate. Hashlimit is used to tune the
            amount of events logged. See the iptables-extensions docs:
            http://ipset.netfilter.org/iptables-extensions.man.html
        hashlimitmode: Set the hashlimit mode, a comma separated string of
            options (srcip,srcport,dstip,dstport). More options here results
            in more logs generated.
        hashlimitburst: Maximum initial number of packets to match.
        hashlimitexpire: Number of milliseconds to keep entries in the hash
            table.
        srcexclude: Exclude a comma separated string of source IPs/CIDRs from
            logging.
    """

    def __init__(self, interface=None, interface_addr=None,
                 log_prefix='"[nfsinkhole] "',
                 protocol='all', dport='0:65535',
                 hashlimit='1/h', hashlimitmode='srcip,dstip,dstport',
                 hashlimitburst='1', hashlimitexpire='3600000',
                 srcexclude='127.0.0.1'
                 ):

        # TODO: add arg checks across all classes
        self.interface = interface
        self.interface_addr = interface_addr
        self.log_prefix = log_prefix
        self.protocol = protocol
        self.dport = dport
        self.hashlimit = hashlimit
        self.hashlimitmode = hashlimitmode
        self.hashlimitburst = hashlimitburst
        self.hashlimitexpire = hashlimitexpire
        self.srcexclude = srcexclude

    def list_existing_rules(self, filter_io_drop=False):
        """
        The function for retrieving current iptables rules related to
        nfsinkhole.

        Args:
            filter_io_drop: Boolean for only showing the DROP rules for INPUT
                and OUTPUT. These are not shown by default. This exists to
                avoid allowing packets on the interface if the service is down.
                If installed, the interface always drops all traffic regardless
                of the service state.

        Returns:
            List: Matching sinkhole lines returned by iptables -S.

        Raises:
            IPTablesError: A Unix process had an error (stderr).
        """

        # Get list summary of iptables rules
        cmd = ['iptables', '-S']

        existing = []

        # Get all of the iptables rules
        try:

            out, err = popen_wrapper(cmd, sudo=True)

        except OSError as e:

            raise IPTablesError('Error encountered when running process "{0}":'
                                '\n{1}'.format(' '.join(cmd), e))

        # If any errors, iterate them and write to log, then raise
        # IPTablesError.
        if err:

            arr = err.splitlines()
            raise IPTablesError('Error encountered when running process "{0}":'
                                '\n{1}'.format(' '.join(cmd), '\n'.join(arr)))

        # Iterate the iptables rules, only grabbing nfsinkhole related rules.
        arr = out.splitlines()
        for line in arr:
            tmp_line = line.decode('ascii', 'ignore')
            if ('SINKHOLE' in tmp_line
                ) or filter_io_drop and tmp_line.strip() in [
                    '-A INPUT -i {0} -j DROP'.format(self.interface),
                    '-A OUTPUT -o {0} -j DROP'.format(self.interface)
            ]:
                existing.append(tmp_line.strip())

        return existing

    def create_rules(self):
        """
        The function for writing iptables rules related to nfsinkhole.
        """

        log.info('Checking for existing iptables rules.')
        existing = self.list_existing_rules()

        # Existing sinkhole related iptables lines found, can't create.
        if len(existing) > 0:

            raise IPTablesExists('Existing iptables rules found for '
                                 'nfsinkhole:\n{0}'
                                 ''.format('\n'.join(existing)))

        log.info('Writing iptables config')

        # Create a new iptables chain for logging
        tmp_arr = ['iptables', '-N', 'SINKHOLE']

        log.info('Writing: {0}'.format(' '.join(tmp_arr)))
        popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        # Exclude IPs/CIDRs from logging (scanners, monitoring, pen-testers,
        # etc):
        for addr in self.srcexclude.split(','):

            tmp_arr = [
                'iptables',
                '-A', 'SINKHOLE',
                '-s', addr,
                '-j', 'RETURN'
            ]

            log.info('Writing: {0}'.format(' '.join(tmp_arr)))
            popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        # Tell the chain to log and use the prefix self.log_prefix:
        tmp_arr = [
            'iptables',
            '-A', 'SINKHOLE',
            '-j', 'LOG',
            '--log-prefix', self.log_prefix
        ]

        log.info('Writing: {0}'.format(' '.join(tmp_arr)))
        popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        # Tell the chain to also log to netfilter (for packet capture):
        tmp_arr = ['iptables', '-A', 'SINKHOLE',
                   '-j', 'NFLOG']

        log.info('Writing: {0}'.format(' '.join(tmp_arr)))
        popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        # Tell the chain to trigger on hashlimit and protocol/port settings
        tmp_arr = [
            'iptables',
            '-i', self.interface,
            '-d', self.interface_addr,
            '-j', 'SINKHOLE',
            '-m', 'hashlimit',
            '--hashlimit', self.hashlimit,
            '--hashlimit-burst', self.hashlimitburst,
            '--hashlimit-mode', self.hashlimitmode,
            '--hashlimit-name', 'sinkhole',
            '--hashlimit-htable-expire', self.hashlimitexpire,
            '-I', 'INPUT', '1'
        ]

        # if --protocol filtered, set mode to multiport with protocol, and
        # set destination port if provided and applicable to the protocol(s)
        if self.protocol != 'all':

            tmp_arr += ['-m', 'multiport', '--protocol', self.protocol]

            if self.dport != '0:65535':
                tmp_arr += ['--dport', self.dport]

        log.info('Writing: {0}'.format(' '.join(tmp_arr)))
        popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

    def create_drop_rule(self):
        """
        The function for writing the iptables DROP rule for the interface.
        """

        log.info('Checking for existing iptables DROP rules.')
        existing = self.list_existing_rules(filter_io_drop=True)

        # Existing sinkhole related iptables lines found, can't create.
        if len(existing) > 0:

            for line in existing:

                if line in (
                    '-A INPUT -i {0} -j DROP'.format(self.interface),
                    '-A OUTPUT -o {0} -j DROP'.format(self.interface)
                ):

                    raise IPTablesExists('Existing iptables DROP rules found '
                                         'for nfsinkhole:\n{0}'
                                         ''.format(line))

        log.info('Writing iptables DROP config')

        # Create rules to drop all I/O traffic:
        tmp_arr = [
            'iptables',
            '-i', self.interface,
            '-j', 'DROP',
            '-I', 'INPUT', '1'
        ]

        log.info('Writing: {0}'.format(' '.join(tmp_arr)))
        popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        tmp_arr = [
            'iptables',
            '-o', self.interface,
            '-j', 'DROP',
            '-I', 'OUTPUT', '1'
        ]

        log.info('Writing: {0}'.format(' '.join(tmp_arr)))

        # TODO: replicate exception handling to other subprocess calls
        try:

            popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        except SubprocessError as e:

            raise IPTablesError(e)

    def delete_rules(self):
        """
        The function for deleting iptables rules related to nfsinkhole.
        """

        log.info('Checking for existing iptables rules.')
        existing = self.list_existing_rules()

        # No sinkhole related iptables lines found.
        if len(existing) == 0:

            raise IPTablesNotExists('No existing rules found.')

        log.info('Deleting iptables config (only what was created)')

        # Iterate all of the active sinkhole related iptables lines
        flush = False
        for line in existing:

            if '-A SINKHOLE' in line or line == '-N SINKHOLE':

                # Don't try to delete the SINKHOLE chain yet, it needs to be
                # empty. Set flush to clear it after this loop.
                flush = True

            elif line not in (
                '-A INPUT -i {0} -j DROP'.format(self.interface),
                '-A OUTPUT -o {0} -j DROP'.format(self.interface)
            ):

                # Delete a single line (not the SINKHOLE chain itself).
                stmt = line.replace('-A', '-D').strip().split(' ')
                tmp_arr = ['iptables'] + stmt

                log.info('Deleting: {0}'.format(' '.join(tmp_arr)))
                popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        # The SINKHOLE chain was detected. Remove it.
        if flush:

            # All lines in the SINKHOLE chain should have been flushed already,
            # but run a flush to be sure.
            tmp_arr = ['iptables', '-F', 'SINKHOLE']

            log.info('Flushing: {0}'.format(' '.join(tmp_arr)))
            popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

            # Now that the SINKHOLE chain has been flushed, we can delete it.
            tmp_arr = ['iptables', '-X', 'SINKHOLE']

            log.info('Deleting: {0}'.format(' '.join(tmp_arr)))
            popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        # Return a list of matching lines.
        return len(existing)

    def delete_drop_rule(self):
        """
        The function for deleting the iptables DROP rule for the interface.
        """

        log.info('Checking for existing iptables DROP rules.')
        existing = self.list_existing_rules(filter_io_drop=True)

        # No sinkhole related iptables lines found.
        if len(existing) == 0:

            raise IPTablesNotExists('No existing rules found.')

        log.info('Deleting iptables DROP config.')

        count = 0
        for line in existing:

            if line in (
                    '-A INPUT -i {0} -j DROP'.format(self.interface),
                    '-A OUTPUT -o {0} -j DROP'.format(self.interface)
            ):
                count += 1
                # Delete a single line (not the SINKHOLE chain itself).
                stmt = line.replace('-A', '-D').split(' ')
                tmp_arr = ['iptables'] + stmt

                log.info('Deleting: {0}'.format(' '.join(tmp_arr)))
                popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        # Return the number of matching lines.
        return count
<commit_msg>coverage exclusions<commit_after># Copyright (c) 2016 Philip Hane
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

from .exceptions import (IPTablesError, IPTablesExists, IPTablesNotExists,
                         SubprocessError)
from .utils import popen_wrapper
import logging

log = logging.getLogger(__name__)

class IPTablesSinkhole:
    """
    The class for managing sinkhole configuration within iptables.

    Args:
        interface: The secondary network interface dedicated to sinkhole
            traffic.
            Warning: Do not accidentally set this to your primary interface.
            It will drop all traffic, and kill your remote access.
        interface_addr: The IP address assigned to interface.
        log_prefix: Prefix for syslog messages.
        protocol: The protocol(s) to log (all traffic will still be dropped).
            Accepts a comma separated string of protocols
            (tcp,udp,udplite,icmp,esp,ah,sctp) or all.
        dport: The destination port(s) to log (for applicable protocols).
            Range should be in the format startport:endport or 0,1,2,3,n..
        hashlimit: Set the hashlimit rate. Hashlimit is used to tune the
            amount of events logged. See the iptables-extensions docs:
            http://ipset.netfilter.org/iptables-extensions.man.html
        hashlimitmode: Set the hashlimit mode, a comma separated string of
            options (srcip,srcport,dstip,dstport). More options here results
            in more logs generated.
        hashlimitburst: Maximum initial number of packets to match.
        hashlimitexpire: Number of milliseconds to keep entries in the hash
            table.
        srcexclude: Exclude a comma separated string of source IPs/CIDRs from
            logging.
    """

    def __init__(self, interface=None, interface_addr=None,
                 log_prefix='"[nfsinkhole] "',
                 protocol='all', dport='0:65535',
                 hashlimit='1/h', hashlimitmode='srcip,dstip,dstport',
                 hashlimitburst='1', hashlimitexpire='3600000',
                 srcexclude='127.0.0.1'
                 ):

        # TODO: add arg checks across all classes
        self.interface = interface
        self.interface_addr = interface_addr
        self.log_prefix = log_prefix
        self.protocol = protocol
        self.dport = dport
        self.hashlimit = hashlimit
        self.hashlimitmode = hashlimitmode
        self.hashlimitburst = hashlimitburst
        self.hashlimitexpire = hashlimitexpire
        self.srcexclude = srcexclude

    def list_existing_rules(self, filter_io_drop=False):
        """
        The function for retrieving current iptables rules related to
        nfsinkhole.

        Args:
            filter_io_drop: Boolean for only showing the DROP rules for INPUT
                and OUTPUT. These are not shown by default. This exists to
                avoid allowing packets on the interface if the service is down.
                If installed, the interface always drops all traffic regardless
                of the service state.

        Returns:
            List: Matching sinkhole lines returned by iptables -S.

        Raises:
            IPTablesError: A Unix process had an error (stderr).
        """

        # Get list summary of iptables rules
        cmd = ['iptables', '-S']

        existing = []

        # Get all of the iptables rules
        try:

            out, err = popen_wrapper(cmd, sudo=True)

        except OSError as e:  # pragma: no cover

            raise IPTablesError('Error encountered when running process "{0}":'
                                '\n{1}'.format(' '.join(cmd), e))

        # If any errors, iterate them and write to log, then raise
        # IPTablesError.
        if err:  # pragma: no cover

            arr = err.splitlines()
            raise IPTablesError('Error encountered when running process "{0}":'
                                '\n{1}'.format(' '.join(cmd), '\n'.join(arr)))

        # Iterate the iptables rules, only grabbing nfsinkhole related rules.
        arr = out.splitlines()
        for line in arr:
            tmp_line = line.decode('ascii', 'ignore')
            if ('SINKHOLE' in tmp_line
                ) or filter_io_drop and tmp_line.strip() in [
                    '-A INPUT -i {0} -j DROP'.format(self.interface),
                    '-A OUTPUT -o {0} -j DROP'.format(self.interface)
            ]:
                existing.append(tmp_line.strip())

        return existing

    def create_rules(self):
        """
        The function for writing iptables rules related to nfsinkhole.
        """

        log.info('Checking for existing iptables rules.')
        existing = self.list_existing_rules()

        # Existing sinkhole related iptables lines found, can't create.
        if len(existing) > 0:

            raise IPTablesExists('Existing iptables rules found for '
                                 'nfsinkhole:\n{0}'
                                 ''.format('\n'.join(existing)))

        log.info('Writing iptables config')

        # Create a new iptables chain for logging
        tmp_arr = ['iptables', '-N', 'SINKHOLE']

        log.info('Writing: {0}'.format(' '.join(tmp_arr)))
        popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        # Exclude IPs/CIDRs from logging (scanners, monitoring, pen-testers,
        # etc):
        for addr in self.srcexclude.split(','):

            tmp_arr = [
                'iptables',
                '-A', 'SINKHOLE',
                '-s', addr,
                '-j', 'RETURN'
            ]

            log.info('Writing: {0}'.format(' '.join(tmp_arr)))
            popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        # Tell the chain to log and use the prefix self.log_prefix:
        tmp_arr = [
            'iptables',
            '-A', 'SINKHOLE',
            '-j', 'LOG',
            '--log-prefix', self.log_prefix
        ]

        log.info('Writing: {0}'.format(' '.join(tmp_arr)))
        popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        # Tell the chain to also log to netfilter (for packet capture):
        tmp_arr = ['iptables', '-A', 'SINKHOLE',
                   '-j', 'NFLOG']

        log.info('Writing: {0}'.format(' '.join(tmp_arr)))
        popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        # Tell the chain to trigger on hashlimit and protocol/port settings
        tmp_arr = [
            'iptables',
            '-i', self.interface,
            '-d', self.interface_addr,
            '-j', 'SINKHOLE',
            '-m', 'hashlimit',
            '--hashlimit', self.hashlimit,
            '--hashlimit-burst', self.hashlimitburst,
            '--hashlimit-mode', self.hashlimitmode,
            '--hashlimit-name', 'sinkhole',
            '--hashlimit-htable-expire', self.hashlimitexpire,
            '-I', 'INPUT', '1'
        ]

        # if --protocol filtered, set mode to multiport with protocol, and
        # set destination port if provided and applicable to the protocol(s)
        if self.protocol != 'all':

            tmp_arr += ['-m', 'multiport', '--protocol', self.protocol]

            if self.dport != '0:65535':
                tmp_arr += ['--dport', self.dport]

        log.info('Writing: {0}'.format(' '.join(tmp_arr)))
        popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

    def create_drop_rule(self):
        """
        The function for writing the iptables DROP rule for the interface.
        """

        log.info('Checking for existing iptables DROP rules.')
        existing = self.list_existing_rules(filter_io_drop=True)

        # Existing sinkhole related iptables lines found, can't create.
        if len(existing) > 0:

            for line in existing:

                if line in (
                    '-A INPUT -i {0} -j DROP'.format(self.interface),
                    '-A OUTPUT -o {0} -j DROP'.format(self.interface)
                ):

                    raise IPTablesExists('Existing iptables DROP rules found '
                                         'for nfsinkhole:\n{0}'
                                         ''.format(line))

        log.info('Writing iptables DROP config')

        # Create rules to drop all I/O traffic:
        tmp_arr = [
            'iptables',
            '-i', self.interface,
            '-j', 'DROP',
            '-I', 'INPUT', '1'
        ]

        log.info('Writing: {0}'.format(' '.join(tmp_arr)))
        popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        tmp_arr = [
            'iptables',
            '-o', self.interface,
            '-j', 'DROP',
            '-I', 'OUTPUT', '1'
        ]

        log.info('Writing: {0}'.format(' '.join(tmp_arr)))

        # TODO: replicate exception handling to other subprocess calls
        try:

            popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        except SubprocessError as e:  # pragma: no cover

            raise IPTablesError(e)

    def delete_rules(self):
        """
        The function for deleting iptables rules related to nfsinkhole.
        """

        log.info('Checking for existing iptables rules.')
        existing = self.list_existing_rules()

        # No sinkhole related iptables lines found.
        if len(existing) == 0:

            raise IPTablesNotExists('No existing rules found.')

        log.info('Deleting iptables config (only what was created)')

        # Iterate all of the active sinkhole related iptables lines
        flush = False
        for line in existing:

            if '-A SINKHOLE' in line or line == '-N SINKHOLE':

                # Don't try to delete the SINKHOLE chain yet, it needs to be
                # empty. Set flush to clear it after this loop.
                flush = True

            elif line not in (
                '-A INPUT -i {0} -j DROP'.format(self.interface),
                '-A OUTPUT -o {0} -j DROP'.format(self.interface)
            ):

                # Delete a single line (not the SINKHOLE chain itself).
                stmt = line.replace('-A', '-D').strip().split(' ')
                tmp_arr = ['iptables'] + stmt

                log.info('Deleting: {0}'.format(' '.join(tmp_arr)))
                popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        # The SINKHOLE chain was detected. Remove it.
        if flush:

            # All lines in the SINKHOLE chain should have been flushed already,
            # but run a flush to be sure.
            tmp_arr = ['iptables', '-F', 'SINKHOLE']

            log.info('Flushing: {0}'.format(' '.join(tmp_arr)))
            popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

            # Now that the SINKHOLE chain has been flushed, we can delete it.
            tmp_arr = ['iptables', '-X', 'SINKHOLE']

            log.info('Deleting: {0}'.format(' '.join(tmp_arr)))
            popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        # Return a list of matching lines.
        return len(existing)

    def delete_drop_rule(self):
        """
        The function for deleting the iptables DROP rule for the interface.
        """

        log.info('Checking for existing iptables DROP rules.')
        existing = self.list_existing_rules(filter_io_drop=True)

        # No sinkhole related iptables lines found.
        if len(existing) == 0:

            raise IPTablesNotExists('No existing rules found.')

        log.info('Deleting iptables DROP config.')

        count = 0
        for line in existing:

            if line in (
                    '-A INPUT -i {0} -j DROP'.format(self.interface),
                    '-A OUTPUT -o {0} -j DROP'.format(self.interface)
            ):
                count += 1
                # Delete a single line (not the SINKHOLE chain itself).
                stmt = line.replace('-A', '-D').split(' ')
                tmp_arr = ['iptables'] + stmt

                log.info('Deleting: {0}'.format(' '.join(tmp_arr)))
                popen_wrapper(cmd_arr=tmp_arr, raise_err=True, sudo=True)

        # Return the number of matching lines.
        return count
<|endoftext|>


<commit_before>"""Tests for disk storage broker."""

import os
import pytz
import datetime

import pytest

from . import tmp_dir_fixture  # NOQA
from . import TEST_SAMPLE_DATA

def test_initialise():

    from dtoolcore.storagebroker import DiskStorageBroker

    path = '/a/path'
    storagebroker = DiskStorageBroker(uri=path)  # NOQA

def test_create_structure(tmp_dir_fixture):  # NOQA

    from dtoolcore.storagebroker import DiskStorageBroker
    from dtoolcore.storagebroker import StorageBrokerOSError

    storagebroker = DiskStorageBroker(tmp_dir_fixture)

    with pytest.raises(StorageBrokerOSError):
        storagebroker.create_structure()

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')

    storagebroker = DiskStorageBroker(destination_path)

    assert not os.path.exists(destination_path)
    storagebroker.create_structure()
    assert os.path.isdir(destination_path)

    destination_path = os.path.join(tmp_dir_fixture, 'sub', 'my_proto_dataset')

    storagebroker = DiskStorageBroker(destination_path)

    with pytest.raises(OSError):
        storagebroker.create_structure()

def test_store_and_retrieve_readme(tmp_dir_fixture):  # NOQA

    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    storagebroker.put_readme('Hello world')
    assert storagebroker.get_readme_content() == 'Hello world'

def test_update_readme(tmp_dir_fixture):  # NOQA

    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    storagebroker.put_readme('Hello world')
    assert storagebroker.get_readme_content() == 'Hello world'

    assert len(storagebroker.list_historical_readme_keys()) == 0

    storagebroker.update_readme('Updated')
    assert storagebroker.get_readme_content() == 'Updated'

    assert len(storagebroker.list_historical_readme_keys()) == 1

    with open(storagebroker.list_historical_readme_keys()[0]) as fh:
        assert fh.read() == 'Hello world'

    storagebroker.update_readme('Updated again')
    assert storagebroker.get_readme_content() == 'Updated again'

    assert len(storagebroker.list_historical_readme_keys()) == 2

def test_store_and_retrieve_admin_metadata(tmp_dir_fixture):  # NOQA

    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    admin_metadata = {'hello': 'world'}
    storagebroker.put_admin_metadata(admin_metadata)

    storagebroker_2 = DiskStorageBroker(destination_path)
    retrieved_admin_metadata = storagebroker_2.get_admin_metadata()
    assert retrieved_admin_metadata == admin_metadata

def test_put_item(tmp_dir_fixture):  # NOQA

    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    input_file_path = os.path.join(TEST_SAMPLE_DATA, 'tiny.png')

    storagebroker.put_item(
        fpath=input_file_path,
        relpath='tiny.png'
    )

    handles = list(storagebroker.iter_item_handles())

    assert 'tiny.png' in handles

def test_store_and_retrieve_manifest(tmp_dir_fixture):  # NOQA
    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    manifest = {'a': 'b', 'c': [1, 2, 3]}

    storagebroker.put_manifest(manifest)

    retrieved_manifest = storagebroker.get_manifest()

    assert retrieved_manifest == manifest

    # Test the formatting on disk.
#   expected = '{\n  "a": "b", \n  "c": [\n    1, \n    2, \n    3\n  ]\n}'
    expected_lines = [
        '{',
        '  "a": "b",',
        '  "c": [',
        '    1,',
        '    2,',
        '    3',
        '  ]',
        '}'
    ]

    with open(storagebroker.get_manifest_key()) as fh:
        for i, actual in enumerate(fh):
            actual = actual.rstrip()
            expected = expected_lines[i]
            assert actual == expected

def test_item_properties(tmp_dir_fixture):  # NOQA
    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    input_file_path = os.path.join(TEST_SAMPLE_DATA, 'tiny.png')

    storagebroker.put_item(
        fpath=input_file_path,
        relpath='tiny.png'
    )

    handles = list(storagebroker.iter_item_handles())

    handle = handles[0]

    item_properties = storagebroker.item_properties(handle)

    # Check size_in_bytes property
    assert item_properties['size_in_bytes'] == 276

    # Check timestamp property
    assert 'utc_timestamp' in item_properties

    time_from_item = datetime.datetime.fromtimestamp(
        float(item_properties['utc_timestamp']),
        tz=pytz.UTC
    )
    time_delta = datetime.datetime.now(tz=pytz.UTC) - time_from_item
    assert time_delta.days == 0
    assert time_delta.seconds < 20

    # Check hash property
    from dtoolcore.filehasher import md5sum_hexdigest
    expected_hash = md5sum_hexdigest(input_file_path)

    assert item_properties['hash'] == expected_hash

    # Check relpath property
    assert item_properties['relpath'] == 'tiny.png'

def test_store_and_retrieve_item_metadata(tmp_dir_fixture):  # NOQA
    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    handle = 'dummy'

    # Here we add two set of metadata with different keys
    storagebroker.add_item_metadata(
        handle=handle,
        key='foo',
        value='bar'
    )
    storagebroker.add_item_metadata(
        handle=handle,
        key='key',
        value={'subkey': 'subval',
               'morekey': 'moreval'}
    )

    # Test metadata retrieval (we get back both sets of metadata)
    metadata = storagebroker.get_item_metadata(handle)
    assert metadata == {
        'foo': 'bar',
        'key': {
            'subkey': 'subval',
            'morekey': 'moreval'
        }
    }

def test_store_and_retrieve_item_metadata(tmp_dir_fixture):  # NOQA
    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    example_overlay = {
        'abcdef': 1,
        'ghijkl': 2
    }

    storagebroker.put_overlay(
        overlay_name="example",
        overlay=example_overlay
    )

    retrieved_overlay = storagebroker.get_overlay('example')

    assert example_overlay == retrieved_overlay

def test_post_freeze_hook(tmp_dir_fixture):  # NOQA
    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    # The below should not raise an OSError because the .dtool/tmp_fragments
    # directory has not been created.
    storagebroker.post_freeze_hook()

    handle = 'dummy'
    storagebroker.add_item_metadata(handle, key='foo', value='bar')

    assert os.path.isdir(storagebroker._metadata_fragments_abspath)
    storagebroker.post_freeze_hook()
    assert not os.path.isdir(storagebroker._metadata_fragments_abspath)

def test_has_admin_metadata(tmp_dir_fixture):  # NOQA

    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    assert not storagebroker.has_admin_metadata()

    storagebroker.create_structure()
    assert not storagebroker.has_admin_metadata()

    admin_metadata = {'hello': 'world'}
    storagebroker.put_admin_metadata(admin_metadata)
    assert storagebroker.has_admin_metadata()

def test_list_dataset_uris(tmp_dir_fixture):  # NOQA

    import dtoolcore
    from dtoolcore.storagebroker import DiskStorageBroker

    assert [] == DiskStorageBroker.list_dataset_uris(
        base_uri=tmp_dir_fixture,
        config_path=None
    )

    # Create two datasets to be copied.
    expected_uris = []
    for name in ["test_ds_1", "test_ds_2"]:
        admin_metadata = dtoolcore.generate_admin_metadata(name)
        proto_dataset = dtoolcore.generate_proto_dataset(
            admin_metadata=admin_metadata,
            base_uri="file://{}".format(tmp_dir_fixture),
            )
        proto_dataset.create()
        expected_uris.append(proto_dataset.uri)

    actual_uris = DiskStorageBroker.list_dataset_uris(
        base_uri=tmp_dir_fixture,
        config_path=None
    )

    assert set(expected_uris) == set(actual_uris)

def test_pre_freeze_hook(tmp_dir_fixture):  # NOQA
    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    # Add a data file.
    data_fpath = os.path.join(storagebroker._data_abspath, "sample.txt")
    with open(data_fpath, "w") as fh:
        fh.write("some sample data")

    # The below should not raise an DiskStorageBrokerValidationError
    # because the structure is correct.
    storagebroker.pre_freeze_hook()

    # Now we add a rogue file.
    rogue_fpath = os.path.join(destination_path, "rogue.txt")
    with open(rogue_fpath, "w") as fh:
        fh.write("I should not be here")

    from dtoolcore.storagebroker import DiskStorageBrokerValidationWarning
    with pytest.raises(DiskStorageBrokerValidationWarning):
        storagebroker.pre_freeze_hook()
<commit_msg>Make the test more robust<commit_after>"""Tests for disk storage broker."""

import os
import pytz
import datetime
import time

import pytest

from . import tmp_dir_fixture  # NOQA
from . import TEST_SAMPLE_DATA

def test_initialise():

    from dtoolcore.storagebroker import DiskStorageBroker

    path = '/a/path'
    storagebroker = DiskStorageBroker(uri=path)  # NOQA

def test_create_structure(tmp_dir_fixture):  # NOQA

    from dtoolcore.storagebroker import DiskStorageBroker
    from dtoolcore.storagebroker import StorageBrokerOSError

    storagebroker = DiskStorageBroker(tmp_dir_fixture)

    with pytest.raises(StorageBrokerOSError):
        storagebroker.create_structure()

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')

    storagebroker = DiskStorageBroker(destination_path)

    assert not os.path.exists(destination_path)
    storagebroker.create_structure()
    assert os.path.isdir(destination_path)

    destination_path = os.path.join(tmp_dir_fixture, 'sub', 'my_proto_dataset')

    storagebroker = DiskStorageBroker(destination_path)

    with pytest.raises(OSError):
        storagebroker.create_structure()

def test_store_and_retrieve_readme(tmp_dir_fixture):  # NOQA

    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    storagebroker.put_readme('Hello world')
    assert storagebroker.get_readme_content() == 'Hello world'

def test_update_readme(tmp_dir_fixture):  # NOQA

    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    storagebroker.put_readme('Hello world')
    assert storagebroker.get_readme_content() == 'Hello world'

    assert len(storagebroker.list_historical_readme_keys()) == 0

    storagebroker.update_readme('Updated')
    assert storagebroker.get_readme_content() == 'Updated'

    assert len(storagebroker.list_historical_readme_keys()) == 1

    with open(storagebroker.list_historical_readme_keys()[0]) as fh:
        assert fh.read() == 'Hello world'

    time.sleep(0.1)

    storagebroker.update_readme('Updated again')
    assert storagebroker.get_readme_content() == 'Updated again'

    assert len(storagebroker.list_historical_readme_keys()) == 2

def test_store_and_retrieve_admin_metadata(tmp_dir_fixture):  # NOQA

    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    admin_metadata = {'hello': 'world'}
    storagebroker.put_admin_metadata(admin_metadata)

    storagebroker_2 = DiskStorageBroker(destination_path)
    retrieved_admin_metadata = storagebroker_2.get_admin_metadata()
    assert retrieved_admin_metadata == admin_metadata

def test_put_item(tmp_dir_fixture):  # NOQA

    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    input_file_path = os.path.join(TEST_SAMPLE_DATA, 'tiny.png')

    storagebroker.put_item(
        fpath=input_file_path,
        relpath='tiny.png'
    )

    handles = list(storagebroker.iter_item_handles())

    assert 'tiny.png' in handles

def test_store_and_retrieve_manifest(tmp_dir_fixture):  # NOQA
    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    manifest = {'a': 'b', 'c': [1, 2, 3]}

    storagebroker.put_manifest(manifest)

    retrieved_manifest = storagebroker.get_manifest()

    assert retrieved_manifest == manifest

    # Test the formatting on disk.
#   expected = '{\n  "a": "b", \n  "c": [\n    1, \n    2, \n    3\n  ]\n}'
    expected_lines = [
        '{',
        '  "a": "b",',
        '  "c": [',
        '    1,',
        '    2,',
        '    3',
        '  ]',
        '}'
    ]

    with open(storagebroker.get_manifest_key()) as fh:
        for i, actual in enumerate(fh):
            actual = actual.rstrip()
            expected = expected_lines[i]
            assert actual == expected

def test_item_properties(tmp_dir_fixture):  # NOQA
    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    input_file_path = os.path.join(TEST_SAMPLE_DATA, 'tiny.png')

    storagebroker.put_item(
        fpath=input_file_path,
        relpath='tiny.png'
    )

    handles = list(storagebroker.iter_item_handles())

    handle = handles[0]

    item_properties = storagebroker.item_properties(handle)

    # Check size_in_bytes property
    assert item_properties['size_in_bytes'] == 276

    # Check timestamp property
    assert 'utc_timestamp' in item_properties

    time_from_item = datetime.datetime.fromtimestamp(
        float(item_properties['utc_timestamp']),
        tz=pytz.UTC
    )
    time_delta = datetime.datetime.now(tz=pytz.UTC) - time_from_item
    assert time_delta.days == 0
    assert time_delta.seconds < 20

    # Check hash property
    from dtoolcore.filehasher import md5sum_hexdigest
    expected_hash = md5sum_hexdigest(input_file_path)

    assert item_properties['hash'] == expected_hash

    # Check relpath property
    assert item_properties['relpath'] == 'tiny.png'

def test_store_and_retrieve_item_metadata(tmp_dir_fixture):  # NOQA
    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    handle = 'dummy'

    # Here we add two set of metadata with different keys
    storagebroker.add_item_metadata(
        handle=handle,
        key='foo',
        value='bar'
    )
    storagebroker.add_item_metadata(
        handle=handle,
        key='key',
        value={'subkey': 'subval',
               'morekey': 'moreval'}
    )

    # Test metadata retrieval (we get back both sets of metadata)
    metadata = storagebroker.get_item_metadata(handle)
    assert metadata == {
        'foo': 'bar',
        'key': {
            'subkey': 'subval',
            'morekey': 'moreval'
        }
    }

def test_store_and_retrieve_item_metadata(tmp_dir_fixture):  # NOQA
    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    example_overlay = {
        'abcdef': 1,
        'ghijkl': 2
    }

    storagebroker.put_overlay(
        overlay_name="example",
        overlay=example_overlay
    )

    retrieved_overlay = storagebroker.get_overlay('example')

    assert example_overlay == retrieved_overlay

def test_post_freeze_hook(tmp_dir_fixture):  # NOQA
    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    # The below should not raise an OSError because the .dtool/tmp_fragments
    # directory has not been created.
    storagebroker.post_freeze_hook()

    handle = 'dummy'
    storagebroker.add_item_metadata(handle, key='foo', value='bar')

    assert os.path.isdir(storagebroker._metadata_fragments_abspath)
    storagebroker.post_freeze_hook()
    assert not os.path.isdir(storagebroker._metadata_fragments_abspath)

def test_has_admin_metadata(tmp_dir_fixture):  # NOQA

    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    assert not storagebroker.has_admin_metadata()

    storagebroker.create_structure()
    assert not storagebroker.has_admin_metadata()

    admin_metadata = {'hello': 'world'}
    storagebroker.put_admin_metadata(admin_metadata)
    assert storagebroker.has_admin_metadata()

def test_list_dataset_uris(tmp_dir_fixture):  # NOQA

    import dtoolcore
    from dtoolcore.storagebroker import DiskStorageBroker

    assert [] == DiskStorageBroker.list_dataset_uris(
        base_uri=tmp_dir_fixture,
        config_path=None
    )

    # Create two datasets to be copied.
    expected_uris = []
    for name in ["test_ds_1", "test_ds_2"]:
        admin_metadata = dtoolcore.generate_admin_metadata(name)
        proto_dataset = dtoolcore.generate_proto_dataset(
            admin_metadata=admin_metadata,
            base_uri="file://{}".format(tmp_dir_fixture),
            )
        proto_dataset.create()
        expected_uris.append(proto_dataset.uri)

    actual_uris = DiskStorageBroker.list_dataset_uris(
        base_uri=tmp_dir_fixture,
        config_path=None
    )

    assert set(expected_uris) == set(actual_uris)

def test_pre_freeze_hook(tmp_dir_fixture):  # NOQA
    from dtoolcore.storagebroker import DiskStorageBroker

    destination_path = os.path.join(tmp_dir_fixture, 'my_proto_dataset')
    storagebroker = DiskStorageBroker(destination_path)

    storagebroker.create_structure()

    # Add a data file.
    data_fpath = os.path.join(storagebroker._data_abspath, "sample.txt")
    with open(data_fpath, "w") as fh:
        fh.write("some sample data")

    # The below should not raise an DiskStorageBrokerValidationError
    # because the structure is correct.
    storagebroker.pre_freeze_hook()

    # Now we add a rogue file.
    rogue_fpath = os.path.join(destination_path, "rogue.txt")
    with open(rogue_fpath, "w") as fh:
        fh.write("I should not be here")

    from dtoolcore.storagebroker import DiskStorageBrokerValidationWarning
    with pytest.raises(DiskStorageBrokerValidationWarning):
        storagebroker.pre_freeze_hook()
<|endoftext|>