<commit_before>"""NISO SUSHI support"""
from __future__ import absolute_import

import datetime
import logging
import uuid

import arrow

from lxml import etree
from lxml import objectify

import requests

import pycounter.constants
import pycounter.exceptions
from pycounter.helpers import convert_date_run
import pycounter.report

logger = logging.getLogger(__name__)
NS = pycounter.constants.NS

def get_sushi_stats_raw(wsdl_url, start_date, end_date, requestor_id=None,
                        requestor_email=None, requestor_name=None,
                        customer_reference=None, customer_name=None,
                        report="JR1", release=4, sushi_dump=False):
    """Get SUSHI stats for a given site in raw XML format.

    :param wsdl_url: URL to SOAP WSDL for this provider

    :param start_date: start date for report (must be first day of a month)

    :param end_date: end date for report (must be last day of a month)

    :param requestor_id: requestor ID as defined by SUSHI protocol

    :param requestor_email: requestor email address, if required by provider

    :param requestor_name: Internationally recognized organization name

    :param customer_reference: customer reference number as defined by SUSHI
        protocol

    :param customer_name: Internationally recognized organization name

    :param report: report type, values defined by SUSHI protocol

    :param release: report release number (should generally be `4`.)

    :param sushi_dump: produces dump of XML to DEBUG logger

    """
    root = etree.Element("{%(SOAP-ENV)s}Envelope" % NS, nsmap=NS)
    body = etree.SubElement(root, "{%(SOAP-ENV)s}Body" % NS)
    timestamp = arrow.utcnow().isoformat()
    rr = etree.SubElement(body, "{%(sushicounter)s}ReportRequest" % NS,
                          {
                              'created': timestamp,
                              'id': str(uuid.uuid4())
                          })

    req = etree.SubElement(rr, "{%(sushi)s}Requestor" % NS)
    rid = etree.SubElement(req, "{%(sushi)s}ID" % NS)
    rid.text = requestor_id
    req_name_element = etree.SubElement(req, "{%(sushi)s}Name" % NS)
    req_name_element.text = requestor_name
    req_email_element = etree.SubElement(req, "{%(sushi)s}Email" % NS)
    req_email_element.text = requestor_email

    cust_ref_elem = etree.SubElement(rr, "{%(sushi)s}CustomerReference" % NS)
    cid = etree.SubElement(cust_ref_elem, "{%(sushi)s}ID" % NS)
    cid.text = customer_reference
    cust_name_elem = etree.SubElement(cust_ref_elem, "{%(sushi)s}Name" % NS)
    cust_name_elem.text = customer_name

    report_def_elem = etree.SubElement(rr, "{%(sushi)s}ReportDefinition" % NS,
                                       Name=report, Release=str(release))
    filters = etree.SubElement(report_def_elem, "{%(sushi)s}Filters" % NS)
    udr = etree.SubElement(filters, "{%(sushi)s}UsageDateRange" % NS)
    beg = etree.SubElement(udr, "{%(sushi)s}Begin" % NS)
    beg.text = start_date.strftime("%Y-%m-%d")
    end = etree.SubElement(udr, "{%(sushi)s}End" % NS)
    end.text = end_date.strftime("%Y-%m-%d")

    payload = etree.tostring(root, pretty_print=True,
                             xml_declaration=True, encoding="utf-8")

    headers = {"SOAPAction": '"SushiService:GetReportIn"',
               "Content-Type": "text/xml; charset=UTF-8",
               "Content-Length": len(payload)}

    response = requests.post(url=wsdl_url,
                             headers=headers,
                             data=payload,
                             verify=False)

    if sushi_dump:
        logger.debug("SUSHI DUMP: request: %s \n\n response: %s",
                     payload,
                     response.content)
    return response.content

def get_report(*args, **kwargs):
    """Get a usage report from a SUSHI server

    returns a :class:`pycounter.report.CounterReport` object.

    parameters: see get_sushi_stats_raw
    """
    raw_report = get_sushi_stats_raw(*args, **kwargs)
    return _raw_to_full(raw_report)

def _ns(namespace, name):
    """Convenience function to make a namespaced XML name.

    :param namespace: one of 'SOAP-ENV', 'sushi', 'sushicounter', 'counter'
    :param name: tag name within the given namespace
    """
    return "{" + NS[namespace] + "}" + name

def _raw_to_full(raw_report):
    """Convert a raw report to a :class:`pycounter.report.CounterReport` object
    """
    try:
        root = etree.fromstring(raw_report)
    except etree.XMLSyntaxError as e:
        logger.error("XML syntax error: %s", raw_report)
        raise pycounter.exceptions.SushiException(e)
    o_root = objectify.fromstring(raw_report)
    rep = None
    try:
        rep = o_root.Body[_ns('sushicounter', "ReportResponse")]
        c_report = rep.Report[_ns('counter', 'Report')]
    except AttributeError:
        try:
            c_report = rep.Report[_ns('counter', 'Reports')].Report
        except AttributeError:
            logger.error("report not found in XML: %s", raw_report)
            raise pycounter.exceptions.SushiException
    logger.debug("COUNTER report: %s", etree.tostring(c_report))
    start_date = datetime.datetime.strptime(
        root.find('.//%s' % _ns('sushi', 'Begin')).text,
        "%Y-%m-%d").date()

    end_date = datetime.datetime.strptime(
        root.find('.//%s' % _ns('sushi', 'End')).text,
        "%Y-%m-%d").date()

    report_data = {'period': (start_date, end_date)}

    rep_def = root.find('.//%s' % _ns('sushi', 'ReportDefinition'))
    report_data['report_version'] = int(rep_def.get('Release'))

    report_data['report_type'] = rep_def.get('Name')

    customer = root.find('.//%s' % _ns('counter', 'Customer'))
    try:
        report_data['customer'] = (customer.find('.//%s' %
                                                 _ns('counter', 'Name')).text)
    except AttributeError:
        report_data['customer'] = ""

    inst_id = customer.find('.//%s' % _ns('counter', 'ID')).text
    report_data['institutional_identifier'] = inst_id

    rep_root = root.find('.//%s' % _ns('counter', 'Report'))
    created_string = rep_root.get('Created')
    if created_string is not None:
        report_data['date_run'] = arrow.get(created_string)
    else:
        report_data['date_run'] = datetime.datetime.now()

    report = pycounter.report.CounterReport(**report_data)

    report.metric = pycounter.constants.METRICS.get(report_data['report_type'])

    for item in c_report.Customer.ReportItems:
        try:
            publisher_name = item.ItemPublisher.text
        except AttributeError:
            publisher_name = ""
        title = item.ItemName.text
        platform = item.ItemPlatform.text

        eissn = issn = isbn = ""

        try:
            for identifier in item.ItemIdentifier:
                if identifier.Type == "Print_ISSN":
                    issn = identifier.Value.text
                    if issn is None:
                        issn = ""
                elif identifier.Type == "Online_ISSN":
                    eissn = identifier.Value.text
                    if eissn is None:
                        eissn = ""
                elif identifier.Type == "Online_ISBN":
                    logging.debug("FOUND ISBN")
                    isbn = identifier.Value.text
                    if isbn is None:
                        isbn = ""

        except AttributeError:
            pass

        month_data = []
        html_usage = 0
        pdf_usage = 0

        for perform_item in item.ItemPerformance:
            item_date = convert_date_run(perform_item.Period.Begin.text)
            logger.debug("perform_item date: %r", item_date)
            usage = None
            for inst in perform_item.Instance:
                if inst.MetricType == "ft_total":
                    usage = str(inst.Count)
                elif inst.MetricType == "ft_pdf":
                    pdf_usage += int(inst.Count)
                elif inst.MetricType == "ft_html":
                    html_usage += int(inst.Count)
            if usage is not None:
                month_data.append((item_date, int(usage)))

        if report.report_type:
            if report.report_type.startswith('JR'):
                report.pubs.append(pycounter.report.CounterJournal(
                    title=title,
                    platform=platform,
                    publisher=publisher_name,
                    period=report.period,
                    metric=report.metric,
                    issn=issn,
                    eissn=eissn,
                    month_data=month_data,
                    html_total=html_usage,
                    pdf_total=pdf_usage
                ))
            elif report.report_type.startswith('BR'):
                report.pubs.append(
                    pycounter.report.CounterBook(
                        title=title,
                        platform=platform,
                        publisher=publisher_name,
                        period=report.period,
                        metric=report.metric,
                        issn=issn,
                        isbn=isbn,
                        month_data=month_data,
                    ))

    return report
<commit_msg>Added support for DB reports to be collected through SUSHI<commit_after>"""NISO SUSHI support"""
from __future__ import absolute_import

import datetime
import logging
import uuid
import collections

import arrow

from lxml import etree
from lxml import objectify

import requests
import six

import pycounter.constants
import pycounter.exceptions
from pycounter.helpers import convert_date_run
import pycounter.report

logger = logging.getLogger(__name__)
NS = pycounter.constants.NS

def get_sushi_stats_raw(wsdl_url, start_date, end_date, requestor_id=None,
                        requestor_email=None, requestor_name=None,
                        customer_reference=None, customer_name=None,
                        report="JR1", release=4, sushi_dump=False):
    """Get SUSHI stats for a given site in raw XML format.

    :param wsdl_url: URL to SOAP WSDL for this provider

    :param start_date: start date for report (must be first day of a month)

    :param end_date: end date for report (must be last day of a month)

    :param requestor_id: requestor ID as defined by SUSHI protocol

    :param requestor_email: requestor email address, if required by provider

    :param requestor_name: Internationally recognized organization name

    :param customer_reference: customer reference number as defined by SUSHI
        protocol

    :param customer_name: Internationally recognized organization name

    :param report: report type, values defined by SUSHI protocol

    :param release: report release number (should generally be `4`.)

    :param sushi_dump: produces dump of XML to DEBUG logger

    """
    root = etree.Element("{%(SOAP-ENV)s}Envelope" % NS, nsmap=NS)
    body = etree.SubElement(root, "{%(SOAP-ENV)s}Body" % NS)
    timestamp = arrow.utcnow().isoformat()
    rr = etree.SubElement(body, "{%(sushicounter)s}ReportRequest" % NS,
                          {
                              'created': timestamp,
                              'id': str(uuid.uuid4())
                          })

    req = etree.SubElement(rr, "{%(sushi)s}Requestor" % NS)
    rid = etree.SubElement(req, "{%(sushi)s}ID" % NS)
    rid.text = requestor_id
    req_name_element = etree.SubElement(req, "{%(sushi)s}Name" % NS)
    req_name_element.text = requestor_name
    req_email_element = etree.SubElement(req, "{%(sushi)s}Email" % NS)
    req_email_element.text = requestor_email

    cust_ref_elem = etree.SubElement(rr, "{%(sushi)s}CustomerReference" % NS)
    cid = etree.SubElement(cust_ref_elem, "{%(sushi)s}ID" % NS)
    cid.text = customer_reference
    cust_name_elem = etree.SubElement(cust_ref_elem, "{%(sushi)s}Name" % NS)
    cust_name_elem.text = customer_name

    report_def_elem = etree.SubElement(rr, "{%(sushi)s}ReportDefinition" % NS,
                                       Name=report, Release=str(release))
    filters = etree.SubElement(report_def_elem, "{%(sushi)s}Filters" % NS)
    udr = etree.SubElement(filters, "{%(sushi)s}UsageDateRange" % NS)
    beg = etree.SubElement(udr, "{%(sushi)s}Begin" % NS)
    beg.text = start_date.strftime("%Y-%m-%d")
    end = etree.SubElement(udr, "{%(sushi)s}End" % NS)
    end.text = end_date.strftime("%Y-%m-%d")

    payload = etree.tostring(root, pretty_print=True,
                             xml_declaration=True, encoding="utf-8")

    headers = {"SOAPAction": '"SushiService:GetReportIn"',
               "Content-Type": "text/xml; charset=UTF-8",
               "Content-Length": len(payload)}

    response = requests.post(url=wsdl_url,
                             headers=headers,
                             data=payload,
                             verify=False)

    if sushi_dump:
        logger.debug("SUSHI DUMP: request: %s \n\n response: %s",
                     payload,
                     response.content)
    return response.content

def get_report(*args, **kwargs):
    """Get a usage report from a SUSHI server

    returns a :class:`pycounter.report.CounterReport` object.

    parameters: see get_sushi_stats_raw
    """
    raw_report = get_sushi_stats_raw(*args, **kwargs)
    return _raw_to_full(raw_report)

def _ns(namespace, name):
    """Convenience function to make a namespaced XML name.

    :param namespace: one of 'SOAP-ENV', 'sushi', 'sushicounter', 'counter'
    :param name: tag name within the given namespace
    """
    return "{" + NS[namespace] + "}" + name

def _raw_to_full(raw_report):
    """Convert a raw report to a :class:`pycounter.report.CounterReport` object
    """
    try:
        root = etree.fromstring(raw_report)
    except etree.XMLSyntaxError as e:
        logger.error("XML syntax error: %s", raw_report)
        raise pycounter.exceptions.SushiException(e)
    o_root = objectify.fromstring(raw_report)
    rep = None
    try:
        rep = o_root.Body[_ns('sushicounter', "ReportResponse")]
        c_report = rep.Report[_ns('counter', 'Report')]
    except AttributeError:
        try:
            c_report = rep.Report[_ns('counter', 'Reports')].Report
        except AttributeError:
            logger.error("report not found in XML: %s", raw_report)
            raise pycounter.exceptions.SushiException
    logger.debug("COUNTER report: %s", etree.tostring(c_report))
    start_date = datetime.datetime.strptime(
        root.find('.//%s' % _ns('sushi', 'Begin')).text,
        "%Y-%m-%d").date()

    end_date = datetime.datetime.strptime(
        root.find('.//%s' % _ns('sushi', 'End')).text,
        "%Y-%m-%d").date()

    report_data = {'period': (start_date, end_date)}

    rep_def = root.find('.//%s' % _ns('sushi', 'ReportDefinition'))
    report_data['report_version'] = int(rep_def.get('Release'))

    report_data['report_type'] = rep_def.get('Name')

    customer = root.find('.//%s' % _ns('counter', 'Customer'))
    try:
        report_data['customer'] = (customer.find('.//%s' %
                                                 _ns('counter', 'Name')).text)
    except AttributeError:
        report_data['customer'] = ""

    inst_id = customer.find('.//%s' % _ns('counter', 'ID')).text
    report_data['institutional_identifier'] = inst_id

    rep_root = root.find('.//%s' % _ns('counter', 'Report'))
    created_string = rep_root.get('Created')
    if created_string is not None:
        report_data['date_run'] = arrow.get(created_string)
    else:
        report_data['date_run'] = datetime.datetime.now()

    report = pycounter.report.CounterReport(**report_data)

    report.metric = pycounter.constants.METRICS.get(report_data['report_type'])

    for item in c_report.Customer.ReportItems:
        try:
            publisher_name = item.ItemPublisher.text
        except AttributeError:
            publisher_name = ""
        title = item.ItemName.text
        platform = item.ItemPlatform.text

        eissn = issn = isbn = ""

        try:
            for identifier in item.ItemIdentifier:
                if identifier.Type == "Print_ISSN":
                    issn = identifier.Value.text
                    if issn is None:
                        issn = ""
                elif identifier.Type == "Online_ISSN":
                    eissn = identifier.Value.text
                    if eissn is None:
                        eissn = ""
                elif identifier.Type == "Online_ISBN":
                    logging.debug("FOUND ISBN")
                    isbn = identifier.Value.text
                    if isbn is None:
                        isbn = ""

        except AttributeError:
            pass

        month_data = []
        html_usage = 0
        pdf_usage = 0

        metrics_for_db = collections.defaultdict(list)

        for perform_item in item.ItemPerformance:
            item_date = convert_date_run(perform_item.Period.Begin.text)
            logger.debug("perform_item date: %r", item_date)
            usage = None
            for inst in perform_item.Instance:
                if inst.MetricType == "ft_total":
                    usage = str(inst.Count)
                elif inst.MetricType == "ft_pdf":
                    pdf_usage += int(inst.Count)
                elif inst.MetricType == "ft_html":
                    html_usage += int(inst.Count)
                elif report.report_type.startswith('DB'):
                    metrics_for_db[inst.MetricType].append((item_date,
                                                            int(inst.Count)))
            if usage is not None:
                month_data.append((item_date, int(usage)))

        if report.report_type:
            if report.report_type.startswith('JR'):
                report.pubs.append(pycounter.report.CounterJournal(
                    title=title,
                    platform=platform,
                    publisher=publisher_name,
                    period=report.period,
                    metric=report.metric,
                    issn=issn,
                    eissn=eissn,
                    month_data=month_data,
                    html_total=html_usage,
                    pdf_total=pdf_usage
                ))
            elif report.report_type.startswith('BR'):
                report.pubs.append(
                    pycounter.report.CounterBook(
                        title=title,
                        platform=platform,
                        publisher=publisher_name,
                        period=report.period,
                        metric=report.metric,
                        issn=issn,
                        isbn=isbn,
                        month_data=month_data,
                    ))
            elif report.report_type.startswith('DB'):
                for metric, month_data in six.iteritems(metrics_for_db):
                    report.pubs.append(
                        pycounter.report.CounterDatabase(
                            title=title,
                            platform=platform,
                            publisher=publisher_name,
                            period=report.period,
                            metric=metric,
                            month_data=month_data
                        ))

    return report
<|endoftext|>


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

import time, inspect, traceback, argparse, json, uuid, requests
from pprint import pprint
from subprocess import Popen, PIPE
from sys import argv, stdout, stderr
from socket import socket, gethostbyname

try:
    import requests
except ImportError:
    stderr.write('ERROR: Python module "requests" not found, please run "pip install requests".\n')
    exit(1)

try:
    import dns.resolver
except ImportError:
    stderr.write('ERROR: Python module "dnspython" not found, please run "pip install dnspython".\n')
    exit(1)

try:
    from termcolor import colored
except ImportError:
    stderr.write('ERROR: Python module "termcolor" not found, please run "pip install termcolor".\n')
    exit(1)

try:
    from OpenSSL.SSL import TLSv1_METHOD, Context, Connection
except ImportError:
    stderr.write('ERROR: Python module "OpenSSL" not found, please run "pip install pyopenssl".\n')
    exit(1)

PROXY_HOST = None
PROXY_PORT = None
BASE_API_URL = 'https://api.digitalocean.com/v2'
DOCKER_IMAGE_SLUG = 'docker'
DEFAULT_FINGERPRINT = ['d1:b6:92:ea:cc:4c:fe:9c:c5:ef:27:ce:33:1f:ba:61']
DEFAULT_REGION_SLUG = 'nyc3'
DEFAULT_MEMORY_SIZE_SLUG = '512mb'
DEFAULT_VCPUS = 1
DEFAULT_DISK_SIZE = 20
DEFAULT_SLEEP = 5
DEFAULT_BRANCH = 'master'

from functools import wraps

DEFAULT_TRIES = 4
DEFAULT_DELAY = 30
DEFAULT_BACKOFF = 2

def retry(ExceptionToCheck, tries=DEFAULT_TRIES, delay=DEFAULT_DELAY, backoff=DEFAULT_BACKOFF, cdata=None):
    '''Retry calling the decorated function using an exponential backoff.

    http://www.saltycrane.com/blog/2009/11/trying-out-retry-decorator-python/
    original from: http://wiki.python.org/moin/PythonDecoratorLibrary#Retry

    :param ExceptionToCheck: the exception to check. may be a tuple of
        exceptions to check
    :type ExceptionToCheck: Exception or tuple
    :param tries: number of times to try (not retry) before giving up
    :type tries: int
    :param delay: initial delay between retries in seconds
    :type delay: int
    :param backoff: backoff multiplier e.g. value of 2 will double the delay
        each retry
    :type backoff: int
    :param logger: logger to use. If None, print
    :type logger: logging.Logger instance
    '''
    def deco_retry(f):
        @wraps(f)
        def f_retry(*args, **kwargs):
            mtries, mdelay = tries, delay
            while mtries > 0:
                try:
                    return f(*args, **kwargs)
                except ExceptionToCheck, e:
                    logger(message='%s, retrying in %d seconds (mtries=%d): %s' % (repr(e), mdelay, mtries, str(cdata)))
                    time.sleep(mdelay)
                    mtries -= 1
                    mdelay *= backoff
            return f(*args, **kwargs)
        return f_retry  # true decorator
    return deco_retry

def logger(message=None):
    print '%s\n' % repr(message)

def get_public_ip():
    resolver = dns.resolver.Resolver()
    resolver.nameservers=[gethostbyname('resolver1.opendns.com')]
    return str(resolver.query('myip.opendns.com', 'A').rrset[0])

def get_regions(s):
    response = s.get('%s/regions' % BASE_API_URL)
    d = json.loads(response.text)
    slugs = []
    for region in d['regions']:
        slugs.append(region['slug'])

    return slugs

def args():
    parser = argparse.ArgumentParser()
    sp = parser.add_subparsers()
    digitalocean = sp.add_parser('digitalocean')
    digitalocean.add_argument('provider', action='store_const', const='digitalocean', help=argparse.SUPPRESS)
    digitalocean.add_argument('--api_token', type=str, required=True, help='DigitalOcean API v2 secret token')
    digitalocean.add_argument('--client_ip', type=str, required=False, default=get_public_ip(), help='client IP to secure Droplet')
    digitalocean.add_argument('--fingerprint', nargs='+', type=str, required=False, default=DEFAULT_FINGERPRINT, help='SSH key fingerprint')
    digitalocean.add_argument('--region', type=str, required=False, default=DEFAULT_REGION_SLUG, help='region to deploy into; use --list_regions for a list')
    digitalocean.add_argument('--branch', type=str, required=False, default=DEFAULT_BRANCH, help='netflix-proxy branch to deploy (default: %s)' % DEFAULT_BRANCH)
    digitalocean.add_argument('--destroy', action='store_true', required=False, help='Destroy droplet on exit')
    digitalocean.add_argument('--list_regions', action='store_true', required=False, help='list all available regions')
    args = parser.parse_args()
    return args

def create_droplet(s, name, cip, fps, region, branch=DEFAULT_BRANCH):
    user_data = '''
#cloud-config

runcmd:
  - git clone -b %s https://github.com/ab77/netflix-proxy /opt/netflix-proxy && cd /opt/netflix-proxy && ./build.sh -c %s
''' % (branch, cip)

    json_data = {'name': name,
                 'region': region,
                 'size': DEFAULT_MEMORY_SIZE_SLUG,
                 'vcpus': DEFAULT_VCPUS,
                 'disk': DEFAULT_DISK_SIZE,
                 'image': DOCKER_IMAGE_SLUG,
                 'ssh_keys': fps,
                 'backups': False,
                 'ipv6': True,
                 'private_networking': False,
                 'user_data': user_data}

    s.headers.update({'Content-Type': 'application/json'})
    post_body = json.dumps(json_data)
    response = s.post('%s/droplets' % BASE_API_URL, data=post_body)
    d = json.loads(response.text)
    pprint(d)

    @retry(AssertionError, cdata='method=%s()' % inspect.stack()[0][3])
    def wait_for_vm_provisioning_completion_retry(action_url):
        response = s.get(action_url)
        d = json.loads(response.text)
        if 'completed' in d['action']['status']:
            print colored(d['action']['status'], 'green')
            assert True
            return d
        else:
            print colored(d['action']['status'], 'red')
            assert False
            return None

    if 'links' not in d:
        return False
    else:
        return wait_for_vm_provisioning_completion_retry(d['links']['actions'][0]['href'])

def destroy_droplet(s, droplet_id):

    @retry(AssertionError)
    def wait_for_vm_deletion_completion_retry(s, droplet_id):
        response = s.delete('%s/droplets/%d' % (BASE_API_URL,
                                                droplet_id))
        if response.__dict__['status_code'] == 204:
            print colored('DELETE /droplets/%d status code %d' % (droplet_id, response.__dict__['status_code']), 'green')
            assert True
            return response.__dict__
        else:
            print colored('DELETE /droplets/%d status code %d' % (droplet_id, response.__dict__['status_code']), 'red')
            assert False
            return None

    return wait_for_vm_deletion_completion_retry(s, droplet_id)

def get_droplet_id_by_name(s, name):
    response = s.get('%s/droplets' % BASE_API_URL)
    d = json.loads(response.text)
    droplet_id = None
    for droplet in d['droplets']:
        if name in droplet['name']:
            droplet_id = droplet['id']

    return droplet_id

def get_droplet_ip_by_name(s, name):
    response = s.get('%s/droplets' % BASE_API_URL)
    d = json.loads(response.text)
    droplet_id = None
    for droplet in d['droplets']:
        if name in droplet['name']:
            droplet_ip = droplet['networks']['v4'][0]['ip_address']

    return droplet_ip

def ssh_run_command(ip, command):
    result = None
    ssh = Popen(['ssh', '-o', 'UserKnownHostsFile=/dev/null', '-o', 'StrictHostKeyChecking=no',
                 '-i', 'id_rsa.travis', 'root@%s' % ip, command],
                shell=False, stdout=PIPE, stderr=PIPE)
    (stdout, stderr) = ssh.communicate()
    print colored('%s: pid = %d, stdout = %s, stderr = %s, rc = %d' % (inspect.stack()[0][3],
                                                                       ssh.pid,
                                                                       stdout.splitlines(),
                                                                       stderr.splitlines(),
                                                                       ssh.returncode), 'grey')
    return dict({'stdout': stdout.splitlines(),
                 'stderr': stderr.splitlines(),
                 'rc': ssh.returncode,
                 'pid': ssh.pid})

@retry(AssertionError, cdata='method=%s()' % inspect.stack()[0][3])
def docker_test_retry(ip):
    stdout = ssh_run_command(ip, 'docker ps')['stdout']
    # https://docs.docker.com/reference/commandline/ps/
    if len(stdout) < 5: # quick and dirty check (5 lines of output = header + containers), needs improvement..
        print colored('%s: stdout = %s, len(stdout) = %d' % (inspect.stack()[0][3],
                                                             stdout,
                                                             len(stdout)), 'red')
        assert False
        return False
    else:
        print colored('%s: stdout = %s, len(stdout) = %d' % (inspect.stack()[0][3],
                                                             stdout,
                                                             len(stdout)), 'green')
        assert True
        return True

def docker_test(ip):
    return docker_test_retry(ip)

def netflix_proxy_test(ip):

    @retry(AssertionError, cdata='method=%s()' % inspect.stack()[0][3])
    def netflix_proxy_test_retry(ip):
        ssh_run_command(ip, 'tail /var/log/cloud-init-output.log')
        rc = ssh_run_command(ip, "grep -E 'Change your DNS to ([0-9]{1,3}[\.]){3}[0-9]{1,3} and start watching Netflix out of region\.' /var/log/cloud-init-output.log")['rc']
        if rc > 0:
            print colored('%s: SSH return code = %s' % (inspect.stack()[0][3], rc), 'red')
            assert False
            return None
        else:
            print colored('%s: SSH return code = %s' % (inspect.stack()[0][3], rc), 'green')
            assert True
            return rc

    return netflix_proxy_test_retry(ip)

def netflix_openssl_test(ip=None, port=443, hostname='netflix.com'):
    """
    Connect to an SNI-enabled server and request a specific hostname
    """

    @retry(Exception, cdata='method=%s()' % inspect.stack()[0][3])
    def netflix_openssl_test_retry(ip):
        client = socket()

        print 'Connecting...',
        stdout.flush()
        client.connect((ip, port))
        print 'connected', client.getpeername()

        client_ssl = Connection(Context(TLSv1_METHOD), client)
        client_ssl.set_connect_state()
        client_ssl.set_tlsext_host_name(hostname)
        client_ssl.do_handshake()
        cert = client_ssl.get_peer_certificate().get_subject()
        cn = [comp for comp in cert.get_components() if comp[0] in ['CN']]
        client_ssl.close()
        print cn
        if hostname in cn[0][1]:
            return True
        else:
            return False

    if not ip: ip = get_public_ip()
    return netflix_openssl_test_retry(ip)

def netflix_test(ip=None, host='www.netflix.com'):

    @retry(Exception, tries=3, delay=10, backoff=2, cdata='method=%s()' % inspect.stack()[0][3])
    def netflix_openssl_test_retry(ip):
        status_code = requests.get('http://%s' % ip, headers={'Host': host}, timeout=10).status_code
        print '%s: status_code=%s' % (host, status_code)
        if not status_code == 200:
            return False
        else:
            return True

    if not ip: ip = get_public_ip()
    return netflix_openssl_test_retry(ip)

def reboot_test(ip):
    stdout = ssh_run_command(ip, 'sudo reboot')['stdout']
    print colored('%s: stdout = %s' % (inspect.stack()[0][3], stdout), 'grey')
    time.sleep(DEFAULT_SLEEP)
    return docker_test_retry(ip)

if __name__ == '__main__':
    arg = args()
    if arg.api_token:
        name = str(uuid.uuid4())
        droplet_id = None
        s = requests.Session()
        if PROXY_HOST and PROXY_PORT:
            s.verify = False
            s.proxies = {'http' : 'http://%s:%s' % (PROXY_HOST, PROXY_PORT),
                         'https': 'https://%s:%s' % (PROXY_HOST, PROXY_PORT)}
        s.headers.update({'Authorization': 'Bearer %s' % arg.api_token})

        if arg.list_regions:
            pprint(get_regions(s))
            exit(0)

        try:
            print colored('Creating Droplet %s...' % name, 'yellow')
            d = create_droplet(s, name, arg.client_ip, arg.fingerprint, arg.region, branch=arg.branch)
            pprint(d)

            droplet_ip = get_droplet_ip_by_name(s, name)
            print colored('Droplet ipaddr = %s...' % droplet_ip, 'cyan')

            print colored('Checking running Docker containers on Droplet with name = %s, ipaddr = %s...' % (name, droplet_ip), 'yellow')
            result = docker_test(droplet_ip)
            if not result: exit(1)

            print colored('Testing netflix-proxy on Droplet with name = %s, ipaddr = %s...' % (name, droplet_ip), 'yellow')
            rc = netflix_proxy_test(droplet_ip)
            if rc > 0: exit(rc)

            print colored('Rebooting Droplet with name = %s, ipaddr = %s...' % (name, droplet_ip), 'yellow')
            result = reboot_test(droplet_ip)
            if not result: exit(1)

            print colored('sniproxy remote test (OpenSSL) on Droplet with name = %s, ipaddr = %s...' % (name, droplet_ip), 'yellow')
            rc = netflix_openssl_test(ip=droplet_ip)
            if not rc: exit(1)

            print colored('sniproxy remote test (HTTP/S) on Droplet with name = %s, ipaddr = %s...' % (name, droplet_ip), 'yellow')
            rc = netflix_test(ip=droplet_ip)
            if not rc: exit(1)

            print colored('Tested, OK..', 'green')
            exit(0)

        except Exception as e:
            print colored(traceback.print_exc(), 'red')
            exit(1)

        finally:
            if arg.destroy:
                droplet_id = get_droplet_id_by_name(s, name)
                if droplet_id:
                    print colored('Destroying Droplet name = %s, id = %s...' % (name, droplet_id), 'yellow')
                    d = destroy_droplet(s, droplet_id)
                    pprint(d)
<commit_msg>Added caching resolver flag<commit_after>#!/usr/bin/env python

import time, inspect, traceback, argparse, json, uuid, requests
from pprint import pprint
from subprocess import Popen, PIPE
from sys import argv, stdout, stderr
from socket import socket, gethostbyname

try:
    import requests
except ImportError:
    stderr.write('ERROR: Python module "requests" not found, please run "pip install requests".\n')
    exit(1)

try:
    import dns.resolver
except ImportError:
    stderr.write('ERROR: Python module "dnspython" not found, please run "pip install dnspython".\n')
    exit(1)

try:
    from termcolor import colored
except ImportError:
    stderr.write('ERROR: Python module "termcolor" not found, please run "pip install termcolor".\n')
    exit(1)

try:
    from OpenSSL.SSL import TLSv1_METHOD, Context, Connection
except ImportError:
    stderr.write('ERROR: Python module "OpenSSL" not found, please run "pip install pyopenssl".\n')
    exit(1)

PROXY_HOST = None
PROXY_PORT = None
BASE_API_URL = 'https://api.digitalocean.com/v2'
DOCKER_IMAGE_SLUG = 'docker'
DEFAULT_FINGERPRINT = ['d1:b6:92:ea:cc:4c:fe:9c:c5:ef:27:ce:33:1f:ba:61']
DEFAULT_REGION_SLUG = 'nyc3'
DEFAULT_MEMORY_SIZE_SLUG = '512mb'
DEFAULT_VCPUS = 1
DEFAULT_DISK_SIZE = 20
DEFAULT_SLEEP = 5
DEFAULT_BRANCH = 'master'

from functools import wraps

DEFAULT_TRIES = 4
DEFAULT_DELAY = 30
DEFAULT_BACKOFF = 2

def retry(ExceptionToCheck, tries=DEFAULT_TRIES, delay=DEFAULT_DELAY, backoff=DEFAULT_BACKOFF, cdata=None):
    '''Retry calling the decorated function using an exponential backoff.

    http://www.saltycrane.com/blog/2009/11/trying-out-retry-decorator-python/
    original from: http://wiki.python.org/moin/PythonDecoratorLibrary#Retry

    :param ExceptionToCheck: the exception to check. may be a tuple of
        exceptions to check
    :type ExceptionToCheck: Exception or tuple
    :param tries: number of times to try (not retry) before giving up
    :type tries: int
    :param delay: initial delay between retries in seconds
    :type delay: int
    :param backoff: backoff multiplier e.g. value of 2 will double the delay
        each retry
    :type backoff: int
    :param logger: logger to use. If None, print
    :type logger: logging.Logger instance
    '''
    def deco_retry(f):
        @wraps(f)
        def f_retry(*args, **kwargs):
            mtries, mdelay = tries, delay
            while mtries > 0:
                try:
                    return f(*args, **kwargs)
                except ExceptionToCheck, e:
                    logger(message='%s, retrying in %d seconds (mtries=%d): %s' % (repr(e), mdelay, mtries, str(cdata)))
                    time.sleep(mdelay)
                    mtries -= 1
                    mdelay *= backoff
            return f(*args, **kwargs)
        return f_retry  # true decorator
    return deco_retry

def logger(message=None):
    print '%s\n' % repr(message)

def get_public_ip():
    resolver = dns.resolver.Resolver()
    resolver.nameservers=[gethostbyname('resolver1.opendns.com')]
    return str(resolver.query('myip.opendns.com', 'A').rrset[0])

def get_regions(s):
    response = s.get('%s/regions' % BASE_API_URL)
    d = json.loads(response.text)
    slugs = []
    for region in d['regions']:
        slugs.append(region['slug'])

    return slugs

def args():
    parser = argparse.ArgumentParser()
    sp = parser.add_subparsers()
    digitalocean = sp.add_parser('digitalocean')
    digitalocean.add_argument('provider', action='store_const', const='digitalocean', help=argparse.SUPPRESS)
    digitalocean.add_argument('--api_token', type=str, required=True, help='DigitalOcean API v2 secret token')
    digitalocean.add_argument('--client_ip', type=str, required=False, default=get_public_ip(), help='client IP to secure Droplet')
    digitalocean.add_argument('--fingerprint', nargs='+', type=str, required=False, default=DEFAULT_FINGERPRINT, help='SSH key fingerprint')
    digitalocean.add_argument('--region', type=str, required=False, default=DEFAULT_REGION_SLUG, help='region to deploy into; use --list_regions for a list')
    digitalocean.add_argument('--branch', type=str, required=False, default=DEFAULT_BRANCH, help='netflix-proxy branch to deploy (default: %s)' % DEFAULT_BRANCH)
    digitalocean.add_argument('--destroy', action='store_true', required=False, help='Destroy droplet on exit')
    digitalocean.add_argument('--list_regions', action='store_true', required=False, help='list all available regions')
    args = parser.parse_args()
    return args

def create_droplet(s, name, cip, fps, region, branch=DEFAULT_BRANCH):
    user_data = '''
#cloud-config

runcmd:
  - git clone -b %s https://github.com/ab77/netflix-proxy /opt/netflix-proxy && cd /opt/netflix-proxy && ./build.sh -c %s -z 1
''' % (branch, cip)

    json_data = {'name': name,
                 'region': region,
                 'size': DEFAULT_MEMORY_SIZE_SLUG,
                 'vcpus': DEFAULT_VCPUS,
                 'disk': DEFAULT_DISK_SIZE,
                 'image': DOCKER_IMAGE_SLUG,
                 'ssh_keys': fps,
                 'backups': False,
                 'ipv6': True,
                 'private_networking': False,
                 'user_data': user_data}

    s.headers.update({'Content-Type': 'application/json'})
    post_body = json.dumps(json_data)
    response = s.post('%s/droplets' % BASE_API_URL, data=post_body)
    d = json.loads(response.text)
    pprint(d)

    @retry(AssertionError, cdata='method=%s()' % inspect.stack()[0][3])
    def wait_for_vm_provisioning_completion_retry(action_url):
        response = s.get(action_url)
        d = json.loads(response.text)
        if 'completed' in d['action']['status']:
            print colored(d['action']['status'], 'green')
            assert True
            return d
        else:
            print colored(d['action']['status'], 'red')
            assert False
            return None

    if 'links' not in d:
        return False
    else:
        return wait_for_vm_provisioning_completion_retry(d['links']['actions'][0]['href'])

def destroy_droplet(s, droplet_id):

    @retry(AssertionError)
    def wait_for_vm_deletion_completion_retry(s, droplet_id):
        response = s.delete('%s/droplets/%d' % (BASE_API_URL,
                                                droplet_id))
        if response.__dict__['status_code'] == 204:
            print colored('DELETE /droplets/%d status code %d' % (droplet_id, response.__dict__['status_code']), 'green')
            assert True
            return response.__dict__
        else:
            print colored('DELETE /droplets/%d status code %d' % (droplet_id, response.__dict__['status_code']), 'red')
            assert False
            return None

    return wait_for_vm_deletion_completion_retry(s, droplet_id)

def get_droplet_id_by_name(s, name):
    response = s.get('%s/droplets' % BASE_API_URL)
    d = json.loads(response.text)
    droplet_id = None
    for droplet in d['droplets']:
        if name in droplet['name']:
            droplet_id = droplet['id']

    return droplet_id

def get_droplet_ip_by_name(s, name):
    response = s.get('%s/droplets' % BASE_API_URL)
    d = json.loads(response.text)
    droplet_id = None
    for droplet in d['droplets']:
        if name in droplet['name']:
            droplet_ip = droplet['networks']['v4'][0]['ip_address']

    return droplet_ip

def ssh_run_command(ip, command):
    result = None
    ssh = Popen(['ssh', '-o', 'UserKnownHostsFile=/dev/null', '-o', 'StrictHostKeyChecking=no',
                 '-i', 'id_rsa.travis', 'root@%s' % ip, command],
                shell=False, stdout=PIPE, stderr=PIPE)
    (stdout, stderr) = ssh.communicate()
    print colored('%s: pid = %d, stdout = %s, stderr = %s, rc = %d' % (inspect.stack()[0][3],
                                                                       ssh.pid,
                                                                       stdout.splitlines(),
                                                                       stderr.splitlines(),
                                                                       ssh.returncode), 'grey')
    return dict({'stdout': stdout.splitlines(),
                 'stderr': stderr.splitlines(),
                 'rc': ssh.returncode,
                 'pid': ssh.pid})

@retry(AssertionError, cdata='method=%s()' % inspect.stack()[0][3])
def docker_test_retry(ip):
    stdout = ssh_run_command(ip, 'docker ps')['stdout']
    # https://docs.docker.com/reference/commandline/ps/
    if len(stdout) < 5: # quick and dirty check (5 lines of output = header + containers), needs improvement..
        print colored('%s: stdout = %s, len(stdout) = %d' % (inspect.stack()[0][3],
                                                             stdout,
                                                             len(stdout)), 'red')
        assert False
        return False
    else:
        print colored('%s: stdout = %s, len(stdout) = %d' % (inspect.stack()[0][3],
                                                             stdout,
                                                             len(stdout)), 'green')
        assert True
        return True

def docker_test(ip):
    return docker_test_retry(ip)

def netflix_proxy_test(ip):

    @retry(AssertionError, cdata='method=%s()' % inspect.stack()[0][3])
    def netflix_proxy_test_retry(ip):
        ssh_run_command(ip, 'tail /var/log/cloud-init-output.log')
        rc = ssh_run_command(ip, "grep -E 'Change your DNS to ([0-9]{1,3}[\.]){3}[0-9]{1,3} and start watching Netflix out of region\.' /var/log/cloud-init-output.log")['rc']
        if rc > 0:
            print colored('%s: SSH return code = %s' % (inspect.stack()[0][3], rc), 'red')
            assert False
            return None
        else:
            print colored('%s: SSH return code = %s' % (inspect.stack()[0][3], rc), 'green')
            assert True
            return rc

    return netflix_proxy_test_retry(ip)

def netflix_openssl_test(ip=None, port=443, hostname='netflix.com'):
    """
    Connect to an SNI-enabled server and request a specific hostname
    """

    @retry(Exception, cdata='method=%s()' % inspect.stack()[0][3])
    def netflix_openssl_test_retry(ip):
        client = socket()

        print 'Connecting...',
        stdout.flush()
        client.connect((ip, port))
        print 'connected', client.getpeername()

        client_ssl = Connection(Context(TLSv1_METHOD), client)
        client_ssl.set_connect_state()
        client_ssl.set_tlsext_host_name(hostname)
        client_ssl.do_handshake()
        cert = client_ssl.get_peer_certificate().get_subject()
        cn = [comp for comp in cert.get_components() if comp[0] in ['CN']]
        client_ssl.close()
        print cn
        if hostname in cn[0][1]:
            return True
        else:
            return False

    if not ip: ip = get_public_ip()
    return netflix_openssl_test_retry(ip)

def netflix_test(ip=None, host='www.netflix.com'):

    @retry(Exception, tries=3, delay=10, backoff=2, cdata='method=%s()' % inspect.stack()[0][3])
    def netflix_openssl_test_retry(ip):
        status_code = requests.get('http://%s' % ip, headers={'Host': host}, timeout=10).status_code
        print '%s: status_code=%s' % (host, status_code)
        if not status_code == 200:
            return False
        else:
            return True

    if not ip: ip = get_public_ip()
    return netflix_openssl_test_retry(ip)

def reboot_test(ip):
    stdout = ssh_run_command(ip, 'sudo reboot')['stdout']
    print colored('%s: stdout = %s' % (inspect.stack()[0][3], stdout), 'grey')
    time.sleep(DEFAULT_SLEEP)
    return docker_test_retry(ip)

if __name__ == '__main__':
    arg = args()
    if arg.api_token:
        name = str(uuid.uuid4())
        droplet_id = None
        s = requests.Session()
        if PROXY_HOST and PROXY_PORT:
            s.verify = False
            s.proxies = {'http' : 'http://%s:%s' % (PROXY_HOST, PROXY_PORT),
                         'https': 'https://%s:%s' % (PROXY_HOST, PROXY_PORT)}
        s.headers.update({'Authorization': 'Bearer %s' % arg.api_token})

        if arg.list_regions:
            pprint(get_regions(s))
            exit(0)

        try:
            print colored('Creating Droplet %s...' % name, 'yellow')
            d = create_droplet(s, name, arg.client_ip, arg.fingerprint, arg.region, branch=arg.branch)
            pprint(d)

            droplet_ip = get_droplet_ip_by_name(s, name)
            print colored('Droplet ipaddr = %s...' % droplet_ip, 'cyan')

            print colored('Checking running Docker containers on Droplet with name = %s, ipaddr = %s...' % (name, droplet_ip), 'yellow')
            result = docker_test(droplet_ip)
            if not result: exit(1)

            print colored('Testing netflix-proxy on Droplet with name = %s, ipaddr = %s...' % (name, droplet_ip), 'yellow')
            rc = netflix_proxy_test(droplet_ip)
            if rc > 0: exit(rc)

            print colored('Rebooting Droplet with name = %s, ipaddr = %s...' % (name, droplet_ip), 'yellow')
            result = reboot_test(droplet_ip)
            if not result: exit(1)

            print colored('sniproxy remote test (OpenSSL) on Droplet with name = %s, ipaddr = %s...' % (name, droplet_ip), 'yellow')
            rc = netflix_openssl_test(ip=droplet_ip)
            if not rc: exit(1)

            print colored('sniproxy remote test (HTTP/S) on Droplet with name = %s, ipaddr = %s...' % (name, droplet_ip), 'yellow')
            rc = netflix_test(ip=droplet_ip)
            if not rc: exit(1)

            print colored('Tested, OK..', 'green')
            exit(0)

        except Exception as e:
            print colored(traceback.print_exc(), 'red')
            exit(1)

        finally:
            if arg.destroy:
                droplet_id = get_droplet_id_by_name(s, name)
                if droplet_id:
                    print colored('Destroying Droplet name = %s, id = %s...' % (name, droplet_id), 'yellow')
                    d = destroy_droplet(s, droplet_id)
                    pprint(d)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# vim: set fileencoding=UTF-8 :

# Copyright (c) 2012, Jonas Häggqvist <rasher@rasher.dk>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
# * 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.
# * Neither the name of the program nor the names of its contributors may be
#   used to endorse or promote products derived from this software without
#   specific prior written permission.
#
# 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 datetime import datetime
from pprint import pprint
from praw import Reddit
from praw.objects import Submission, Comment
from rules import RuleHandler
import argparse
import logging
import logging.config
import os.path
import praw.errors
import re
import sys
import time

NAME = "ModBot"
VERSION = 0.1

SEEN = set()
SEEN_FILE = 'seen.list'

MODQUEUE_ACTED = set()
MODQUEUE_ACTED_FILE = 'modqueue_acted.list'

def performaction(thing, action, rule, matches):
    origaction = action.strip()
    action = action.strip().lower()
    logging.info("Perform %s on %s from %s" % (action, thing.permalink,
        rule['_filename']))

    # Compose message. All these are made the same way
    if action in ('respond', 'messagemods', 'messageauthor'):
        subject = "Modbot rule matched"
        try:
            if 'subject' in rule:
                subject = rule['subject']
            text = rule['content'].format(thing=thing, rule=rule,
                    matches=matches)
        except Exception:
            # We'll just use a default message then
            text = """
The following post/comment by /u/{thing.author.comment} matched the rule
{rule._filename}: {thing.permalink}
""".strip()
            text = text.format(thing=thing, rule=rule, matches=matches)
            pass

    if action == 'upvote':
        thing.upvote()
    elif action == 'log':
        rulelogger = logging.getLogger('rulelogger')
        rulelogger.info("%s - %s", thing.permalink, rule['_filename'])
    elif action == 'spam':
        thing.remove(True)
    elif action == 'remove':
        thing.remove(False)
    elif action == 'approve':
        thing.approve()
    elif action == 'respond':
        if isinstance(thing, Submission):
            comment = thing.add_comment(text)
        else:
            comment = thing.reply(text)
        comment.distinguish()
    elif action == 'messagemods':
        thing.subreddit.send_message(subject, text)
    elif action in ('beep', 'bell'):
        sys.stdout.write("\x07")
    elif action == 'messageauthor':
        thing.author.send_message(subject, text)
    elif action == 'report':
        thing.report()
    elif action.startswith('linkflair') and ':' in action:
        linkflair = origaction.split(':', 1)[1]
        values = linkflair.split(':', 1)
        if len(values) == 1:
            s = thing.set_flair(flair_text=values[0])
        else:
            css_class, text = values
            s = thing.set_flair(flair_text=text, flair_css_class=css_class)
    elif action in ('none', 'null', 'ignore'):
        pass
    else:
        logging.warning("Unknown action: %s from %s" % (action,
            rule['_filename']))

def applyrule(thing, rule, matches):
    if 'action' in rule:
        for action in rule['action'].split(','):
            performaction(thing, action, rule, matches)
    if 'actions' in rule:
        for action in rule['actions'].split(','):
            performaction(thing, action, rule, matches)

def matchrules(thing, rules, is_modqueue=False):
    print "%s" % is_modqueue
    if thing.name in SEEN and not is_modqueue:
        return False
    if thing.name in MODQUEUE_ACTED and is_modqueue:
        return False
    rulekey = {
            'username': (None, lambda(t): t.author.name),
            'numreports': (None, lambda(t): t.num_reports),
            'domain': (Submission, lambda(t): t.domain),
            'title': (Submission, lambda(t): t.title),
            'upvotes': (None, lambda(t): t.ups),
            'downvotes': (None, lambda(t): t.downs),
            'score': (None, lambda(t): t.score),
            'type': (None,
                lambda(t): 'comment' if isinstance(t, Comment) else
                'submission'),
            'body': (None,
                lambda(t): t.body if isinstance(t, Comment) else t.selftext),
            'bodylength': (None,
                lambda(t): len(t.body) if isinstance(t, Comment) else
                len(t.selftext)),
            'dayhour': (None,
                lambda(t):
                    datetime.fromtimestamp(t.timestamp).strftime("%a-%H")),
            }

    for rule in rules:
        logging.debug("Match %s against %s" % (thing.name, rule['_filename']))
        ruleMatches = True
        matches = {}
        for key, value in rule.iteritems():
            if key not in rulekey:
                continue
            kind, getter = rulekey[key]
            if kind != None and not isinstance(thing, kind):
                ruleMatches = False
                break
            logging.debug("Match %s %s %s" % (thing.name, key,
                unicode(getter(thing))))
            regex = '(?P<full>%s)' % value
            m = re.search(regex, unicode(getter(thing)), flags=re.IGNORECASE)
            if not m:
                ruleMatches = False
                break
            else:
                matches[key] = m.groupdict()
        if ruleMatches:
            try:
                applyrule(thing, rule, matches)
                seen(thing.name)
                if is_modqueue:
                    modqueue_acted(thing.name)
                return True
            except Exception, e:
                logging.error(str(e))
                return False
    seen(thing.name)
    return False

def modqueue_acted(thing_id):
    remember_thing(MODQUEUE_ACTED, MODQUEUE_ACTED_FILE, thing_id)

def seen(thing_id):
    remember_thing(SEEN, SEEN_FILE, thing_id)

def remember_thing(idset, filename, thing_id):
    if thing_id in idset:
        return
    idset.add(thing_id)
    f = open(filename, 'a')
    f.write("%s,%d\n" % (thing_id, int(datetime.utcnow().strftime("%s"))))
    f.close()

def read_thinglists():
    if os.path.exists(SEEN_FILE):
        for line in open(SEEN_FILE):
            SEEN.add(line.strip().split(",")[0])
    if os.path.exists(MODQUEUE_ACTED_FILE):
        for line in open(MODQUEUE_ACTED_FILE):
            MODQUEUE_ACTED.add(line.strip().split(",")[0])

def main():
    parser = argparse.ArgumentParser(
            prog='modbot.py',
            description="Subreddit automoderating script")
    parser.add_argument('subreddit')
    parser.add_argument('-r', '--rulesdir', default='./rules')
    parser.add_argument('-u', '--user')
    parser.add_argument('-p', '--password')
    args = parser.parse_args()

    rh = RuleHandler(args.rulesdir, '*.rule')

    logging.config.fileConfig('logging.conf')

    reddit = Reddit('%s/%s' % (NAME, VERSION))
    try:
        reddit.login(args.user, args.password)
    except praw.errors.InvalidUserPass:
        logging.critical("Login failure")
        sys.exit(1)

    sub = reddit.get_subreddit(args.subreddit)
    read_thinglists()
    comments_ph = None
    submissions_ph = None

    while True:
        logging.info("Loop start")
        rh.update()

        try:
            modqueue_items = sub.get_modqueue(limit=100)
        except Exception, e:
            modqueue_items = []

        num = 0
        for modqueue_item in modqueue_items:
            num += 1
            matchrules(modqueue_item, rh.rules, is_modqueue=True)
        logging.info("Checked %d modqueue items" % num)

        try:
            if comments_ph == None:
                comments = sub.get_comments(limit=10)
            else:
                comments = sub.get_comments(place_holder=comments_ph,
                        limit=100)
        except Exception, e:
            comments = []

        num = 0
        for comment in comments:
            logging.debug("Checking %s start" % comment.name)
            num += 1
            if comments_ph == None or num == 1:
                comments_ph = comment.id
            matchrules(comment, rh.rules)
            logging.debug("Checking %s done" % comment.name)
        logging.info("Checked %d comments" % num)

        try:
            if submissions_ph == None:
                submissions = sub.get_new_by_date(limit=10)
            else:
                submissions = sub.get_new_by_date(place_holder=submissions_ph,
                        limit=100)
        except Exception, e:
            submissions = []

        num = 0
        for submission in submissions:
            num += 1
            if submissions_ph == None or num == 1:
                submissions_ph = submission.id
            matchrules(submission, rh.rules)
        logging.info("Checked %d submissions" % num)

        sys.exit(0)
        time.sleep(30)

if __name__ == "__main__":
    main()
<commit_msg>Move value getting into a seperate class<commit_after>#!/usr/bin/env python
# vim: set fileencoding=UTF-8 :

# Copyright (c) 2012, Jonas Häggqvist <rasher@rasher.dk>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
# * 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.
# * Neither the name of the program nor the names of its contributors may be
#   used to endorse or promote products derived from this software without
#   specific prior written permission.
#
# 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 datetime import datetime
from pprint import pprint
from praw import Reddit
from praw.objects import Submission, Comment, Redditor
from rules import RuleHandler
from decorators import RequiresType
import argparse
import logging
import logging.config
import os.path
import praw.errors
import re
import sys
import time

NAME = "ModBot"
VERSION = 0.1

SEEN = set()
SEEN_FILE = 'seen.list'

MODQUEUE_ACTED = set()
MODQUEUE_ACTED_FILE = 'modqueue_acted.list'

def performaction(thing, action, rule, matches):
    origaction = action.strip()
    action = action.strip().lower()
    logging.info("Perform %s on %s from %s" % (action, thing.permalink,
        rule['_filename']))

    # Compose message. All these are made the same way
    if action in ('respond', 'messagemods', 'messageauthor'):
        subject = "Modbot rule matched"
        try:
            if 'subject' in rule:
                subject = rule['subject']
            text = rule['content'].format(thing=thing, rule=rule,
                    matches=matches)
        except Exception:
            # We'll just use a default message then
            text = """
The following post/comment by /u/{thing.author.comment} matched the rule
{rule._filename}: {thing.permalink}
""".strip()
            text = text.format(thing=thing, rule=rule, matches=matches)
            pass

    if action == 'upvote':
        thing.upvote()
    elif action == 'log':
        rulelogger = logging.getLogger('rulelogger')
        rulelogger.info("%s - %s", thing.permalink, rule['_filename'])
    elif action == 'spam':
        thing.remove(True)
    elif action == 'remove':
        thing.remove(False)
    elif action == 'approve':
        thing.approve()
    elif action == 'respond':
        if isinstance(thing, Submission):
            comment = thing.add_comment(text)
        else:
            comment = thing.reply(text)
        comment.distinguish()
    elif action == 'messagemods':
        thing.subreddit.send_message(subject, text)
    elif action in ('beep', 'bell'):
        sys.stdout.write("\x07")
    elif action == 'messageauthor':
        thing.author.send_message(subject, text)
    elif action == 'report':
        thing.report()
    elif action.startswith('linkflair') and ':' in action:
        linkflair = origaction.split(':', 1)[1]
        values = linkflair.split(':', 1)
        if len(values) == 1:
            s = thing.set_flair(flair_text=values[0])
        else:
            css_class, text = values
            s = thing.set_flair(flair_text=text, flair_css_class=css_class)
    elif action in ('none', 'null', 'ignore'):
        pass
    else:
        logging.warning("Unknown action: %s from %s" % (action,
            rule['_filename']))

def applyrule(thing, rule, matches):
    if 'action' in rule:
        for action in rule['action'].split(','):
            performaction(thing, action, rule, matches)
    if 'actions' in rule:
        for action in rule['actions'].split(','):
            performaction(thing, action, rule, matches)

class ValueGetter:
    """
    Simplify getting a value from a thing, using a common name for different
    types of things, regardless of how the value is retrieved.
    """
    @RequiresType(Submission, Comment, Redditor, position=2)
    def username(self, thing):
        if not isinstance(thing, Redditor):
            thing = thing.author
        return thing.name

    @RequiresType(Submission, Comment, position=2)
    def numreports(self, thing):
        return thing.num_reports

    @RequiresType(Submission, position=2)
    def domain(self, thing):
        return thing.domain

    @RequiresType(Submission, position=2)
    def title(self, thing):
        return thing.title

    @RequiresType(Submission, Comment, position=2)
    def upvotes(self, thing):
        return thing.ups

    @RequiresType(Submission, Comment, position=2)
    def downvotes(self, thing):
        return thing.downs

    @RequiresType(Submission, Comment, position=2)
    def score(self, thing):
        return thing.score

    def type(self, thing):
        return type(thing).__name__.lower()

    @RequiresType(Submission, Comment, position=2)
    def body(self, thing):
        if isinstance(thing, Comment):
            return t.body
        else:
            return t.selftext

    def bodylength(self, thing):
        return len(self.body(thing))

    @RequiresType(Submission, Comment, position=2)
    def dayhour(self, thing):
        return datetime.fromtimestamp(thing.timestamp).strftime("%a-%H")

def matchrules(thing, rules, is_modqueue=False):
    if thing.name in SEEN and not is_modqueue:
        return False
    if thing.name in MODQUEUE_ACTED and is_modqueue:
        return False

    vg = ValueGetter()
    for rule in rules:
        logging.debug("Match %s against %s" % (thing.name, rule['_filename']))
        ruleMatches = True
        matches = {}
        # Ideally we should loop through these in a more clever order
        # e.g. check the type value before checking more expensive ones
        for key, value in rule.iteritems():
            try:
                fieldvalue = unicode(getattr(vg, key)(thing))
            except AttributeError:
                continue
            except TypeError:
                ruleMatches = False
                break

            logging.debug("Match %s %s %s" % (thing.name, key, fieldvalue))
            regex = '(?P<full>%s)' % value
            m = re.search(regex, fieldvalue, flags=re.IGNORECASE)
            if not m:
                ruleMatches = False
                break
            else:
                matches[key] = m.groupdict()
        if ruleMatches:
            try:
                applyrule(thing, rule, matches)
                seen(thing.name)
                if is_modqueue:
                    modqueue_acted(thing.name)
                return True
            except Exception, e:
                logging.error(str(e))
                return False
    seen(thing.name)
    return False

def modqueue_acted(thing_id):
    remember_thing(MODQUEUE_ACTED, MODQUEUE_ACTED_FILE, thing_id)

def seen(thing_id):
    remember_thing(SEEN, SEEN_FILE, thing_id)

def remember_thing(idset, filename, thing_id):
    if thing_id in idset:
        return
    idset.add(thing_id)
    f = open(filename, 'a')
    f.write("%s,%d\n" % (thing_id, int(datetime.utcnow().strftime("%s"))))
    f.close()

def read_thinglists():
    if os.path.exists(SEEN_FILE):
        for line in open(SEEN_FILE):
            SEEN.add(line.strip().split(",")[0])
    if os.path.exists(MODQUEUE_ACTED_FILE):
        for line in open(MODQUEUE_ACTED_FILE):
            MODQUEUE_ACTED.add(line.strip().split(",")[0])

def main():
    parser = argparse.ArgumentParser(
            prog='modbot.py',
            description="Subreddit automoderating script")
    parser.add_argument('subreddit')
    parser.add_argument('-r', '--rulesdir', default='./rules')
    parser.add_argument('-u', '--user')
    parser.add_argument('-p', '--password')
    args = parser.parse_args()

    rh = RuleHandler(args.rulesdir, '*.rule')

    logging.config.fileConfig('logging.conf')

    reddit = Reddit('%s/%s' % (NAME, VERSION))
    try:
        reddit.login(args.user, args.password)
    except praw.errors.InvalidUserPass:
        logging.critical("Login failure")
        sys.exit(1)

    sub = reddit.get_subreddit(args.subreddit)
    read_thinglists()
    comments_ph = None
    submissions_ph = None

    while True:
        logging.info("Loop start")
        rh.update()

        try:
            modqueue_items = sub.get_modqueue(limit=100)
        except Exception, e:
            modqueue_items = []

        num = 0
        for modqueue_item in modqueue_items:
            num += 1
            matchrules(modqueue_item, rh.rules, is_modqueue=True)
        logging.info("Checked %d modqueue items" % num)

        try:
            if comments_ph == None:
                comments = sub.get_comments(limit=10)
            else:
                comments = sub.get_comments(place_holder=comments_ph,
                        limit=100)
        except Exception, e:
            comments = []

        num = 0
        for comment in comments:
            logging.debug("Checking %s start" % comment.name)
            num += 1
            if comments_ph == None or num == 1:
                comments_ph = comment.id
            matchrules(comment, rh.rules)
            logging.debug("Checking %s done" % comment.name)
        logging.info("Checked %d comments" % num)

        try:
            if submissions_ph == None:
                submissions = sub.get_new_by_date(limit=10)
            else:
                submissions = sub.get_new_by_date(place_holder=submissions_ph,
                        limit=100)
        except Exception, e:
            submissions = []

        num = 0
        for submission in submissions:
            num += 1
            if submissions_ph == None or num == 1:
                submissions_ph = submission.id
            matchrules(submission, rh.rules)
        logging.info("Checked %d submissions" % num)

        sys.exit(0)
        time.sleep(30)

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