<commit_before>"""
A collection of functions which test the most basic operations of various services.
"""

import datetime
import logging
import re
import time
import uuid
import urllib3
from io import BytesIO

from django.conf import settings
from django.contrib.auth.models import User
from django.core import cache
from django.db import connections
from django.db.utils import OperationalError

import attr
import gevent
import requests
from celery import Celery

from soil import heartbeat

from corehq.apps.app_manager.models import Application
from corehq.apps.change_feed.connection import get_kafka_client
from corehq.apps.es import GroupES
from corehq.apps.formplayer_api.utils import get_formplayer_url
from corehq.apps.hqadmin.escheck import check_es_cluster_health
from corehq.apps.hqadmin.utils import parse_celery_pings, parse_celery_workers
from corehq.blobs import CODES, get_blob_db
from corehq.celery_monitoring.heartbeat import (
    Heartbeat,
    HeartbeatNeverRecorded,
)
from corehq.elastic import refresh_elasticsearch_index, send_to_elasticsearch
from corehq.util.decorators import change_log_level, ignore_warning
from corehq.util.timer import TimingContext

@attr.s
class ServiceStatus(object):
    success = attr.ib()
    msg = attr.ib()
    exception = attr.ib(default=None)
    duration = attr.ib(default=None)

class UnknownCheckException(Exception):
    pass

def check_redis():
    if 'redis' in settings.CACHES:
        import redis
        rc = cache.caches['redis']
        redis_api = redis.StrictRedis.from_url('%s' % rc._server)
        memory = redis_api.info()['used_memory_human']
        result = rc.set('serverup_check_key', 'test', timeout=5)
        return ServiceStatus(result, "Redis is up and using {} memory".format(memory))
    else:
        return ServiceStatus(False, "Redis is not configured on this system!")

def check_all_rabbitmq():
    unwell_rabbits = []
    ip_regex = re.compile(r'[0-9]+.[0-9]+.[0-9]+.[0-9]+')

    for broker_url in settings.CELERY_BROKER_URL.split(';'):
        check_status, failure = check_rabbitmq(broker_url)
        if not check_status:
            failed_rabbit_ip = ip_regex.search(broker_url).group()
            unwell_rabbits.append((failed_rabbit_ip, failure))

    if not unwell_rabbits:
        return ServiceStatus(True, 'RabbitMQ OK')

    else:
        return ServiceStatus(False, '; '.join(['{}:{}'.format(rabbit[0], rabbit[1])
                                        for rabbit in unwell_rabbits])
                      )

def check_rabbitmq(broker_url):
    if broker_url.startswith('amqp'):
        amqp_parts = broker_url.replace('amqp://', '').split('/')
        mq_management_url = amqp_parts[0].replace('5672', '15672')
        vhost = amqp_parts[1]
        try:
            vhost_dict = requests.get('http://%s/api/vhosts' % mq_management_url, timeout=2).json()
            for d in vhost_dict:
                if d['name'] == vhost:
                    return True, 'RabbitMQ OK'
            return False, 'RabbitMQ Offline'
        except Exception as e:
            return False, "RabbitMQ Error: %s" % e
    elif settings.CELERY_BROKER_URL.startswith('redis'):
        return True, "RabbitMQ Not configured, but not needed"
    else:
        return False, "RabbitMQ Not configured"

@change_log_level('kafka.client', logging.WARNING)
def check_kafka():
    try:
        client = get_kafka_client()
    except Exception as e:
        return ServiceStatus(False, "Could not connect to Kafka: %s" % e)

    if len(client.cluster.brokers()) == 0:
        return ServiceStatus(False, "No Kafka brokers found")
    elif len(client.cluster.topics()) == 0:
        return ServiceStatus(False, "No Kafka topics found")
    else:
        return ServiceStatus(True, "Kafka seems to be in order")

@change_log_level('urllib3.connectionpool', logging.WARNING)
def check_elasticsearch():
    cluster_health = check_es_cluster_health()
    if cluster_health == 'red':
        return ServiceStatus(False, "Cluster health at %s" % cluster_health)

    doc = {'_id': 'elasticsearch-service-check-{}'.format(uuid.uuid4().hex[:7]),
           'date': datetime.datetime.now().isoformat()}
    try:
        send_to_elasticsearch('groups', doc)
        refresh_elasticsearch_index('groups')
        hits = GroupES().remove_default_filters().doc_id(doc['_id']).run().hits
        if doc in hits:
            return ServiceStatus(True, "Successfully sent a doc to ES and read it back")
        else:
            return ServiceStatus(False, "Something went wrong sending a doc to ES")
    finally:
        send_to_elasticsearch('groups', doc, delete=True)  # clean up

def check_blobdb():
    """Save something to the blobdb and try reading it back."""
    db = get_blob_db()
    contents = b"It takes Pluto 248 Earth years to complete one orbit!"
    meta = db.put(
        BytesIO(contents),
        domain="<unknown>",
        parent_id="check_blobdb",
        type_code=CODES.tempfile,
    )
    with db.get(key=meta.key) as fh:
        res = fh.read()
    db.delete(key=meta.key)
    if res == contents:
        return ServiceStatus(True, "Successfully saved a file to the blobdb")
    return ServiceStatus(False, "Failed to save a file to the blobdb")

def check_celery():
    blocked_queues = []

    for queue, threshold in settings.CELERY_HEARTBEAT_THRESHOLDS.items():
        if threshold:
            threshold = datetime.timedelta(seconds=threshold)
            try:
                blockage_duration = Heartbeat(queue).get_and_report_blockage_duration()
            except HeartbeatNeverRecorded:
                blocked_queues.append((queue, 'as long as we can see', threshold))
            else:
                # We get a lot of self-resolving celery "downtime" under 5 minutes
                # so to make actionable, we never alert on blockage under 5 minutes
                # It is still counted as out of SLA for the celery uptime metric in datadog
                if blockage_duration > max(threshold, datetime.timedelta(minutes=5)):
                    blocked_queues.append((queue, blockage_duration, threshold))

    if blocked_queues:
        return ServiceStatus(False, '\n'.join(
            "{} has been blocked for {} (max allowed is {})".format(
                queue, blockage_duration, threshold
            ) for queue, blockage_duration, threshold in blocked_queues))
    else:
        return ServiceStatus(True, "OK")

def check_heartbeat():
    is_alive = heartbeat.is_alive()
    return ServiceStatus(is_alive, "OK" if is_alive else "DOWN")

def check_postgres():
    connected = True
    status_str = ""
    for db in settings.DATABASES:
        db_conn = connections[db]
        try:
            c = db_conn.cursor()
            c_status = 'OK'
        except OperationalError:
            c_status = 'FAIL'
            connected = False
        status_str += "%s:%s:%s " % (db, settings.DATABASES[db]['NAME'], c_status)

    a_user = User.objects.first()
    if a_user is None:
        status_str += "No users found in postgres"
    else:
        status_str += "Successfully got a user from postgres"

    if a_user is None or not connected:
        return ServiceStatus(False, status_str)
    return ServiceStatus(True, status_str)

def check_couch():
    """Confirm CouchDB is up and running, by hitting an arbitrary view."""
    results = Application.view('app_manager/builds_by_date', limit=1).all()
    assert isinstance(results, list), "Couch didn't return a list of builds"
    return ServiceStatus(True, "Successfully queried an arbitrary couch view")

@ignore_warning(urllib3.exceptions.InsecureRequestWarning)
def check_formplayer():
    try:
        # Setting verify=False in this request keeps this from failing for urls with self-signed certificates.
        # Allowing this because the certificate will always be self-signed in the "provable deploy"
        # bootstrapping test in commcare-cloud.
        res = requests.get('{}/serverup'.format(get_formplayer_url()), timeout=5, verify=False)
    except requests.exceptions.ConnectTimeout:
        return ServiceStatus(False, "Could not establish a connection in time")
    except requests.ConnectionError:
        return ServiceStatus(False, "Could not connect to formplayer")
    else:
        msg = "Formplayer returned a {} status code".format(res.status_code)
        return ServiceStatus(res.ok, msg)

def run_checks(checks_to_do):
    greenlets = []
    with TimingContext() as timer:
        for check_name in checks_to_do:
            if check_name not in CHECKS:
                raise UnknownCheckException(check_name)

            greenlets.append(gevent.spawn(_run_check, check_name, timer))
        gevent.joinall(greenlets)
    return [greenlet.value for greenlet in greenlets]

def _run_check(check_name, timer):
    check_info = CHECKS[check_name]
    with timer(check_name):
        try:
            status = check_info['check_func']()
        except Exception as e:
            status = ServiceStatus(False, "{} raised an error".format(check_name), e)
        status.duration = timer.peek().duration
    return check_name, status

CHECKS = {
    'kafka': {
        "always_check": True,
        "check_func": check_kafka,
    },
    'redis': {
        "always_check": True,
        "check_func": check_redis,
    },
    'postgres': {
        "always_check": True,
        "check_func": check_postgres,
    },
    'couch': {
        "always_check": True,
        "check_func": check_couch,
    },
    'celery': {
        "always_check": False,
        "check_func": check_celery,
    },
    'heartbeat': {
        "always_check": False,
        "check_func": check_heartbeat,
    },
    'elasticsearch': {
        "always_check": True,
        "check_func": check_elasticsearch,
    },
    'blobdb': {
        "always_check": True,
        "check_func": check_blobdb,
    },
    'formplayer': {
        "always_check": True,
        "check_func": check_formplayer,
    },
    'rabbitmq': {
        "always_check": True,
        "check_func": check_all_rabbitmq,
    },
}
<commit_msg>include formplayer URL in check response<commit_after>"""
A collection of functions which test the most basic operations of various services.
"""

import datetime
import logging
import re
import time
import uuid
import urllib3
from io import BytesIO

from django.conf import settings
from django.contrib.auth.models import User
from django.core import cache
from django.db import connections
from django.db.utils import OperationalError

import attr
import gevent
import requests
from celery import Celery

from soil import heartbeat

from corehq.apps.app_manager.models import Application
from corehq.apps.change_feed.connection import get_kafka_client
from corehq.apps.es import GroupES
from corehq.apps.formplayer_api.utils import get_formplayer_url
from corehq.apps.hqadmin.escheck import check_es_cluster_health
from corehq.apps.hqadmin.utils import parse_celery_pings, parse_celery_workers
from corehq.blobs import CODES, get_blob_db
from corehq.celery_monitoring.heartbeat import (
    Heartbeat,
    HeartbeatNeverRecorded,
)
from corehq.elastic import refresh_elasticsearch_index, send_to_elasticsearch
from corehq.util.decorators import change_log_level, ignore_warning
from corehq.util.timer import TimingContext

@attr.s
class ServiceStatus(object):
    success = attr.ib()
    msg = attr.ib()
    exception = attr.ib(default=None)
    duration = attr.ib(default=None)

class UnknownCheckException(Exception):
    pass

def check_redis():
    if 'redis' in settings.CACHES:
        import redis
        rc = cache.caches['redis']
        redis_api = redis.StrictRedis.from_url('%s' % rc._server)
        memory = redis_api.info()['used_memory_human']
        result = rc.set('serverup_check_key', 'test', timeout=5)
        return ServiceStatus(result, "Redis is up and using {} memory".format(memory))
    else:
        return ServiceStatus(False, "Redis is not configured on this system!")

def check_all_rabbitmq():
    unwell_rabbits = []
    ip_regex = re.compile(r'[0-9]+.[0-9]+.[0-9]+.[0-9]+')

    for broker_url in settings.CELERY_BROKER_URL.split(';'):
        check_status, failure = check_rabbitmq(broker_url)
        if not check_status:
            failed_rabbit_ip = ip_regex.search(broker_url).group()
            unwell_rabbits.append((failed_rabbit_ip, failure))

    if not unwell_rabbits:
        return ServiceStatus(True, 'RabbitMQ OK')

    else:
        return ServiceStatus(False, '; '.join(['{}:{}'.format(rabbit[0], rabbit[1])
                                        for rabbit in unwell_rabbits])
                      )

def check_rabbitmq(broker_url):
    if broker_url.startswith('amqp'):
        amqp_parts = broker_url.replace('amqp://', '').split('/')
        mq_management_url = amqp_parts[0].replace('5672', '15672')
        vhost = amqp_parts[1]
        try:
            vhost_dict = requests.get('http://%s/api/vhosts' % mq_management_url, timeout=2).json()
            for d in vhost_dict:
                if d['name'] == vhost:
                    return True, 'RabbitMQ OK'
            return False, 'RabbitMQ Offline'
        except Exception as e:
            return False, "RabbitMQ Error: %s" % e
    elif settings.CELERY_BROKER_URL.startswith('redis'):
        return True, "RabbitMQ Not configured, but not needed"
    else:
        return False, "RabbitMQ Not configured"

@change_log_level('kafka.client', logging.WARNING)
def check_kafka():
    try:
        client = get_kafka_client()
    except Exception as e:
        return ServiceStatus(False, "Could not connect to Kafka: %s" % e)

    if len(client.cluster.brokers()) == 0:
        return ServiceStatus(False, "No Kafka brokers found")
    elif len(client.cluster.topics()) == 0:
        return ServiceStatus(False, "No Kafka topics found")
    else:
        return ServiceStatus(True, "Kafka seems to be in order")

@change_log_level('urllib3.connectionpool', logging.WARNING)
def check_elasticsearch():
    cluster_health = check_es_cluster_health()
    if cluster_health == 'red':
        return ServiceStatus(False, "Cluster health at %s" % cluster_health)

    doc = {'_id': 'elasticsearch-service-check-{}'.format(uuid.uuid4().hex[:7]),
           'date': datetime.datetime.now().isoformat()}
    try:
        send_to_elasticsearch('groups', doc)
        refresh_elasticsearch_index('groups')
        hits = GroupES().remove_default_filters().doc_id(doc['_id']).run().hits
        if doc in hits:
            return ServiceStatus(True, "Successfully sent a doc to ES and read it back")
        else:
            return ServiceStatus(False, "Something went wrong sending a doc to ES")
    finally:
        send_to_elasticsearch('groups', doc, delete=True)  # clean up

def check_blobdb():
    """Save something to the blobdb and try reading it back."""
    db = get_blob_db()
    contents = b"It takes Pluto 248 Earth years to complete one orbit!"
    meta = db.put(
        BytesIO(contents),
        domain="<unknown>",
        parent_id="check_blobdb",
        type_code=CODES.tempfile,
    )
    with db.get(key=meta.key) as fh:
        res = fh.read()
    db.delete(key=meta.key)
    if res == contents:
        return ServiceStatus(True, "Successfully saved a file to the blobdb")
    return ServiceStatus(False, "Failed to save a file to the blobdb")

def check_celery():
    blocked_queues = []

    for queue, threshold in settings.CELERY_HEARTBEAT_THRESHOLDS.items():
        if threshold:
            threshold = datetime.timedelta(seconds=threshold)
            try:
                blockage_duration = Heartbeat(queue).get_and_report_blockage_duration()
            except HeartbeatNeverRecorded:
                blocked_queues.append((queue, 'as long as we can see', threshold))
            else:
                # We get a lot of self-resolving celery "downtime" under 5 minutes
                # so to make actionable, we never alert on blockage under 5 minutes
                # It is still counted as out of SLA for the celery uptime metric in datadog
                if blockage_duration > max(threshold, datetime.timedelta(minutes=5)):
                    blocked_queues.append((queue, blockage_duration, threshold))

    if blocked_queues:
        return ServiceStatus(False, '\n'.join(
            "{} has been blocked for {} (max allowed is {})".format(
                queue, blockage_duration, threshold
            ) for queue, blockage_duration, threshold in blocked_queues))
    else:
        return ServiceStatus(True, "OK")

def check_heartbeat():
    is_alive = heartbeat.is_alive()
    return ServiceStatus(is_alive, "OK" if is_alive else "DOWN")

def check_postgres():
    connected = True
    status_str = ""
    for db in settings.DATABASES:
        db_conn = connections[db]
        try:
            c = db_conn.cursor()
            c_status = 'OK'
        except OperationalError:
            c_status = 'FAIL'
            connected = False
        status_str += "%s:%s:%s " % (db, settings.DATABASES[db]['NAME'], c_status)

    a_user = User.objects.first()
    if a_user is None:
        status_str += "No users found in postgres"
    else:
        status_str += "Successfully got a user from postgres"

    if a_user is None or not connected:
        return ServiceStatus(False, status_str)
    return ServiceStatus(True, status_str)

def check_couch():
    """Confirm CouchDB is up and running, by hitting an arbitrary view."""
    results = Application.view('app_manager/builds_by_date', limit=1).all()
    assert isinstance(results, list), "Couch didn't return a list of builds"
    return ServiceStatus(True, "Successfully queried an arbitrary couch view")

@ignore_warning(urllib3.exceptions.InsecureRequestWarning)
def check_formplayer():
    url = f'{get_formplayer_url()}/serverup'
    try:
        # Setting verify=False in this request keeps this from failing for urls with self-signed certificates.
        # Allowing this because the certificate will always be self-signed in the "provable deploy"
        # bootstrapping test in commcare-cloud.
        res = requests.get(url, timeout=5, verify=False)
    except requests.exceptions.ConnectTimeout:
        return ServiceStatus(False, f"Could not establish a connection in time {url}")
    except requests.ConnectionError:
        return ServiceStatus(False, f"Could not connect to formplayer: {url}")
    else:
        msg = f"Formplayer returned a {res.status_code} status code: {url}"
        return ServiceStatus(res.ok, msg)

def run_checks(checks_to_do):
    greenlets = []
    with TimingContext() as timer:
        for check_name in checks_to_do:
            if check_name not in CHECKS:
                raise UnknownCheckException(check_name)

            greenlets.append(gevent.spawn(_run_check, check_name, timer))
        gevent.joinall(greenlets)
    return [greenlet.value for greenlet in greenlets]

def _run_check(check_name, timer):
    check_info = CHECKS[check_name]
    with timer(check_name):
        try:
            status = check_info['check_func']()
        except Exception as e:
            status = ServiceStatus(False, "{} raised an error".format(check_name), e)
        status.duration = timer.peek().duration
    return check_name, status

CHECKS = {
    'kafka': {
        "always_check": True,
        "check_func": check_kafka,
    },
    'redis': {
        "always_check": True,
        "check_func": check_redis,
    },
    'postgres': {
        "always_check": True,
        "check_func": check_postgres,
    },
    'couch': {
        "always_check": True,
        "check_func": check_couch,
    },
    'celery': {
        "always_check": False,
        "check_func": check_celery,
    },
    'heartbeat': {
        "always_check": False,
        "check_func": check_heartbeat,
    },
    'elasticsearch': {
        "always_check": True,
        "check_func": check_elasticsearch,
    },
    'blobdb': {
        "always_check": True,
        "check_func": check_blobdb,
    },
    'formplayer': {
        "always_check": True,
        "check_func": check_formplayer,
    },
    'rabbitmq': {
        "always_check": True,
        "check_func": check_all_rabbitmq,
    },
}
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import unicode_literals

import ghdiff
from memoized import memoized
from six.moves import zip

from corehq.apps.translations.app_translations import (
    expected_bulk_app_sheet_headers,
    expected_bulk_app_sheet_rows,
    get_unicode_dicts,
)
from corehq.apps.translations.const import MODULES_AND_FORMS_SHEET_NAME
from corehq.apps.translations.generators import SKIP_TRANSFEX_STRING, AppTranslationsGenerator

COLUMNS_TO_COMPARE = {
    'module_and_form': ['Type', 'sheet_name'],
    'module': ['case_property', 'list_or_detail'],
    'form': ['label'],
}
# return from ghdiff in case of no differences
NO_GHDIFF_MESSAGE = ghdiff.diff([], [], css=False)

class UploadedTranslationsValidator(object):
    """
    this compares the excel sheet uploaded with translations with what would be generated
    with current app state and flags any discrepancies found between the two
    """
    def __init__(self, app, uploaded_workbook, lang_prefix='default_'):
        self.app = app
        self.uploaded_workbook = uploaded_workbook
        self.headers = None
        self.expected_rows = None
        self.lang_prefix = lang_prefix
        self.default_language_column = self.lang_prefix + self.app.default_language
        self.app_translation_generator = AppTranslationsGenerator(
            self.app.domain, self.app.get_id, None, self.app.default_language, self.app.default_language,
            self.lang_prefix)

    def _generate_expected_headers_and_rows(self):
        self.headers = {h[0]: h[1] for h in expected_bulk_app_sheet_headers(self.app)}
        self.expected_rows = expected_bulk_app_sheet_rows(
            self.app,
            exclude_module=lambda module: SKIP_TRANSFEX_STRING in module.comment,
            exclude_form=lambda form: SKIP_TRANSFEX_STRING in form.comment
        )

    @memoized
    def _get_header_index(self, sheet_name, header):
        for index, _column_name in enumerate(self.headers[sheet_name]):
            if _column_name == header:
                return index

    def _filter_rows(self, for_type, expected_rows, sheet_name):
        if for_type == 'form':
            return self.app_translation_generator._filter_invalid_rows_for_form(
                expected_rows,
                self.app_translation_generator.sheet_name_to_module_or_form_type_and_id[sheet_name].id,
                self._get_header_index(sheet_name, 'label')
            )
        elif for_type == 'module':
            return self.app_translation_generator._filter_invalid_rows_for_module(
                expected_rows,
                self.app_translation_generator.sheet_name_to_module_or_form_type_and_id[sheet_name].id,
                self._get_header_index(sheet_name, 'case_property'),
                self._get_header_index(sheet_name, 'list_or_detail'),
                self._get_header_index(sheet_name, self.default_language_column)
            )
        elif for_type == 'module_and_form':
            return expected_rows
        assert False, "Unexpected type"

    def _compare_sheet(self, sheet_name, uploaded_rows, for_type):
        """
        :param uploaded_rows: dict
        :param for_type: type of sheet, module_and_forms, module, form
        :return: diff generated by ghdiff or None
        """
        columns_to_compare = COLUMNS_TO_COMPARE[for_type] + [self.default_language_column]
        expected_rows = self._filter_rows(for_type, self.expected_rows[sheet_name], sheet_name)

        iterate_on = [expected_rows, uploaded_rows]
        parsed_expected_rows = []
        parsed_uploaded_rows = []
        for i, (expected_row, uploaded_row) in enumerate(zip(*iterate_on), 2):
            parsed_expected_row = [uploaded_row.get(column_name) for column_name in columns_to_compare]
            parsed_uploaded_row = [expected_row[self._get_header_index(sheet_name, column_name)]
                                   for column_name in columns_to_compare]
            parsed_expected_rows.append(parsed_expected_row)
            parsed_uploaded_rows.append(parsed_uploaded_row)
        expected_rows_as_string = '\n'.join([', '.join(row) for row in parsed_expected_rows])
        uploaded_rows_as_string = '\n'.join([', '.join(row) for row in parsed_uploaded_rows])
        diff = ghdiff.diff(expected_rows_as_string, uploaded_rows_as_string, css=False)
        if diff == NO_GHDIFF_MESSAGE:
            return None
        return diff

    def compare(self):
        msgs = {}
        self._generate_expected_headers_and_rows()
        for sheet in self.uploaded_workbook.worksheets:
            sheet_name = sheet.worksheet.title
            if sheet_name not in self.expected_rows:
                continue
            rows = get_unicode_dicts(sheet)
            if sheet_name == MODULES_AND_FORMS_SHEET_NAME:
                error_msgs = self._compare_sheet(sheet_name, rows, 'module_and_form')
            elif 'module' in sheet_name and 'form' not in sheet_name:
                error_msgs = self._compare_sheet(sheet_name, rows, 'module')
            elif 'module' in sheet_name and 'form' in sheet_name:
                error_msgs = self._compare_sheet(sheet_name, rows, 'form')
            else:
                raise Exception("Got unexpected sheet name %s" % sheet_name)
            if error_msgs:
                msgs[sheet_name] = error_msgs
        return msgs
<commit_msg>add comment [skip ci]<commit_after>from __future__ import absolute_import
from __future__ import unicode_literals

import ghdiff
from memoized import memoized
from six.moves import zip

from corehq.apps.translations.app_translations import (
    expected_bulk_app_sheet_headers,
    expected_bulk_app_sheet_rows,
    get_unicode_dicts,
)
from corehq.apps.translations.const import MODULES_AND_FORMS_SHEET_NAME
from corehq.apps.translations.generators import SKIP_TRANSFEX_STRING, AppTranslationsGenerator

COLUMNS_TO_COMPARE = {
    'module_and_form': ['Type', 'sheet_name'],
    'module': ['case_property', 'list_or_detail'],
    'form': ['label'],
}
# return from ghdiff in case of no differences
NO_GHDIFF_MESSAGE = ghdiff.diff([], [], css=False)

class UploadedTranslationsValidator(object):
    """
    this compares the excel sheet uploaded with translations with what would be generated
    with current app state and flags any discrepancies found between the two
    """
    def __init__(self, app, uploaded_workbook, lang_prefix='default_'):
        self.app = app
        self.uploaded_workbook = uploaded_workbook
        self.headers = None
        self.expected_rows = None
        self.lang_prefix = lang_prefix
        self.default_language_column = self.lang_prefix + self.app.default_language
        self.app_translation_generator = AppTranslationsGenerator(
            self.app.domain, self.app.get_id, None, self.app.default_language, self.app.default_language,
            self.lang_prefix)

    def _generate_expected_headers_and_rows(self):
        self.headers = {h[0]: h[1] for h in expected_bulk_app_sheet_headers(self.app)}
        self.expected_rows = expected_bulk_app_sheet_rows(
            self.app,
            exclude_module=lambda module: SKIP_TRANSFEX_STRING in module.comment,
            exclude_form=lambda form: SKIP_TRANSFEX_STRING in form.comment
        )

    @memoized
    def _get_header_index(self, sheet_name, header):
        for index, _column_name in enumerate(self.headers[sheet_name]):
            if _column_name == header:
                return index

    def _filter_rows(self, for_type, expected_rows, sheet_name):
        if for_type == 'form':
            return self.app_translation_generator._filter_invalid_rows_for_form(
                expected_rows,
                self.app_translation_generator.sheet_name_to_module_or_form_type_and_id[sheet_name].id,
                self._get_header_index(sheet_name, 'label')
            )
        elif for_type == 'module':
            return self.app_translation_generator._filter_invalid_rows_for_module(
                expected_rows,
                self.app_translation_generator.sheet_name_to_module_or_form_type_and_id[sheet_name].id,
                self._get_header_index(sheet_name, 'case_property'),
                self._get_header_index(sheet_name, 'list_or_detail'),
                self._get_header_index(sheet_name, self.default_language_column)
            )
        elif for_type == 'module_and_form':
            return expected_rows
        assert False, "Unexpected type"

    def _compare_sheet(self, sheet_name, uploaded_rows, for_type):
        """
        :param uploaded_rows: dict
        :param for_type: type of sheet, module_and_forms, module, form
        :return: diff generated by ghdiff or None
        """
        columns_to_compare = COLUMNS_TO_COMPARE[for_type] + [self.default_language_column]
        expected_rows = self._filter_rows(for_type, self.expected_rows[sheet_name], sheet_name)

        iterate_on = [expected_rows, uploaded_rows]
        parsed_expected_rows = []
        parsed_uploaded_rows = []
        for i, (expected_row, uploaded_row) in enumerate(zip(*iterate_on), 2):
            parsed_expected_row = [uploaded_row.get(column_name) for column_name in columns_to_compare]
            parsed_uploaded_row = [expected_row[self._get_header_index(sheet_name, column_name)]
                                   for column_name in columns_to_compare]
            parsed_expected_rows.append(parsed_expected_row)
            parsed_uploaded_rows.append(parsed_uploaded_row)
        expected_rows_as_string = '\n'.join([', '.join(row) for row in parsed_expected_rows])
        uploaded_rows_as_string = '\n'.join([', '.join(row) for row in parsed_uploaded_rows])
        diff = ghdiff.diff(expected_rows_as_string, uploaded_rows_as_string, css=False)
        if diff == NO_GHDIFF_MESSAGE:
            return None
        return diff

    def compare(self):
        msgs = {}
        self._generate_expected_headers_and_rows()
        for sheet in self.uploaded_workbook.worksheets:
            sheet_name = sheet.worksheet.title
            # if sheet is not in the expected rows, ignore it. This can happen if the module/form sheet is excluded
            # from transifex integration
            if sheet_name not in self.expected_rows:
                continue
            rows = get_unicode_dicts(sheet)
            if sheet_name == MODULES_AND_FORMS_SHEET_NAME:
                error_msgs = self._compare_sheet(sheet_name, rows, 'module_and_form')
            elif 'module' in sheet_name and 'form' not in sheet_name:
                error_msgs = self._compare_sheet(sheet_name, rows, 'module')
            elif 'module' in sheet_name and 'form' in sheet_name:
                error_msgs = self._compare_sheet(sheet_name, rows, 'form')
            else:
                raise Exception("Got unexpected sheet name %s" % sheet_name)
            if error_msgs:
                msgs[sheet_name] = error_msgs
        return msgs
<|endoftext|>


<commit_before>import os
import simplejson as json
import zc.recipe.egg

class CodeintelRecipe(object):

    SUPPORTED_LANGUAGES = ['PHP', 'Python', 'RHTML', 'JavaScript', 'Smarty', 'Mason', 'Node.js', 'XBL', 'Tcl', 'HTML',
        'HTML5', 'TemplateToolkit', 'XUL', 'Django', 'Perl', 'Ruby', 'Python3']

    def __init__(self, buildout, name, options):

        self.buildout, self.name, self.options = buildout, name, options

        # Read .codeintel directory from options or just dump in the buildout root
        default_target = os.path.join(self.buildout['buildout']['directory'], '.codeintel')
        target = options.get("codeintel-path", default_target)

        self.codeintel_directory = target

        if not os.path.exists(self.codeintel_directory):
            os.mkdir(self.codeintel_directory)

        self._python = self.buildout['buildout']['executable']

        # form extra paths dictionary where each key is lang and value is a list of paths
        config = {}
        for lang in self.SUPPORTED_LANGUAGES:
            lang_lower = lang.lower()
            lang_extra_paths = filter(None, map(lambda x: x.strip(),
                options.get('{0}-extra-paths'.format(lang_lower), '').split()))
            if lang_extra_paths:
                config.setdefault(lang, {})
                config[lang]['{0}ExtraPaths'.format(lang_lower)] = lang_extra_paths
            lang_path = options.get('{0}-path'.format(lang_lower), '')
            if lang_path:
                config[lang][lang_lower] = lang_path

        # update formed dict with extra-paths option
        python_extra_paths = options.get('extra-paths', '').split('\n')
        python_extra_paths = filter(lambda p: p.strip() != '', python_extra_paths)

        config.setdefault('Python', {}).setdefault('pythonExtraPaths', []).extend(python_extra_paths)

        self.config = config

        self.egg = zc.recipe.egg.Egg(self.buildout, self.name, self.options)

    def install(self):
        python_paths = (filter(None, map(lambda p: p.strip(), self.egg.working_set()[1].entries)) +
            self.config['Python']['pythonExtraPaths'])
        self.config['Python']['pythonExtraPaths'] = python_paths
        open(os.path.join(self.codeintel_directory, "config"), "w").write(json.dumps(self.config, indent=4,
            sort_keys=True))
        return ""

    def update(self):
        return self.install()
<commit_msg>Minimal python 3.3 support.<commit_after>import os
import simplejson as json
import zc.recipe.egg

class CodeintelRecipe(object):

    SUPPORTED_LANGUAGES = ['PHP', 'Python', 'RHTML', 'JavaScript', 'Smarty', 'Mason', 'Node.js', 'XBL', 'Tcl', 'HTML',
        'HTML5', 'TemplateToolkit', 'XUL', 'Django', 'Perl', 'Ruby', 'Python3']

    def __init__(self, buildout, name, options):

        self.buildout, self.name, self.options = buildout, name, options

        # Read .codeintel directory from options or just dump in the buildout root
        default_target = os.path.join(self.buildout['buildout']['directory'], '.codeintel')
        target = options.get("codeintel-path", default_target)

        self.codeintel_directory = target

        if not os.path.exists(self.codeintel_directory):
            os.mkdir(self.codeintel_directory)

        self._python = self.buildout['buildout']['executable']

        # form extra paths dictionary where each key is lang and value is a list of paths
        config = {}
        for lang in self.SUPPORTED_LANGUAGES:
            lang_lower = lang.lower()
            lang_extra_paths = filter(None, map(lambda x: x.strip(),
                options.get('{0}-extra-paths'.format(lang_lower), '').split()))
            lang_extra_paths = list(lang_extra_paths)
            if lang_extra_paths:
                config.setdefault(lang, {})
                config[lang]['{0}ExtraPaths'.format(lang_lower)] = lang_extra_paths
            lang_path = options.get('{0}-path'.format(lang_lower), '')
            if lang_path:
                config[lang][lang_lower] = lang_path

        # update formed dict with extra-paths option
        python_extra_paths = options.get('extra-paths', '').split('\n')
        python_extra_paths = filter(lambda p: p.strip() != '', python_extra_paths)

        config.setdefault('Python', {}).setdefault('pythonExtraPaths', []).extend(python_extra_paths)

        self.config = config

        self.egg = zc.recipe.egg.Egg(self.buildout, self.name, self.options)

    def install(self):
        python_paths = (
            list(filter(None, map(lambda p: p.strip(), self.egg.working_set()[1].entries))) +
            self.config['Python']['pythonExtraPaths']
        )
        self.config['Python']['pythonExtraPaths'] = python_paths
        open(os.path.join(self.codeintel_directory, "config"), "w").write(json.dumps(self.config, indent=4,
            sort_keys=True))
        return ""

    def update(self):
        return self.install()
<|endoftext|>


<commit_before>#! /usr/bin/env python
from __future__ import print_function, absolute_import, division
'''
Re-request and process COS data to keep products up to date.
'''

import argparse
import yaml
import os
import pdb

from .find_new_cos_data import janky_connect
from .request_data import run_all_retrievals

#-----------------------------------------------------------------------------#
#-----------------------------------------------------------------------------#
def query_proposal(proposal):
    '''
    Find all member datasets of a given proposal ID.

    Parameters:
    -----------
        proposal : int
            The proposal ID for which you wish to retrieve all data.

    Returns:
    --------
        prop_datasets : list
            List of all datasets given a proposal ID.
    '''

    # Open the configuration file for the MAST database connection (TSQL).
    config_file = os.path.join(os.environ['HOME'], "configure2.yaml")
    with open(config_file, 'r') as f:
        SETTINGS = yaml.load(f)

    # Need to use opus_rep DB instead of dadsops_rep to get ALL datasets
    SETTINGS["database"] = "opus_rep"
    query0 = "SELECT DISTINCT program_id FROM executed "\
    "WHERE proposal_id='{0}'\ngo".format(str(proposal))
    vid = janky_connect(SETTINGS, query0)

    # Now we need to go back to dadsop_rep to get ads_data_set_name
    SETTINGS["database"] = "dadsops_rep"
    query1 = "SELECT DISTINCT ads_data_set_name FROM "\
    "archive_data_set_all WHERE ads_data_set_name "\
    "LIKE '_{0}%'\ngo".format(str(vid[0]))
    datasets = janky_connect(SETTINGS, query1)
    prop_datasets = [[x, proposal] for x in datasets]

    return prop_datasets

#-----------------------------------------------------------------------------#
#-----------------------------------------------------------------------------#
def handle_datasets(rootname):
    '''
    Get proposal IDs and, if applicable, all member datasets given a visit or
    program.

    Parameters:
    -----------
        rootname : str
            The rootname (partial or full) to be retrieved.

    Returns:
    --------
        root_datasets : list
            List of all datasets for a given rootname.
    '''

    config_file = os.path.join(os.environ['HOME'], "configure2.yaml")
    with open(config_file, 'r') as f:
        SETTINGS = yaml.load(f)

    query0 = "SELECT DISTINCT ads_data_set_name,ads_pep_id FROM "\
    "archive_data_set_all WHERE ads_data_set_name "\
    "LIKE '{0}%'\ngo".format(str(rootname))
    root_datasets = janky_connect(SETTINGS, query0)

    # If data do not have a PID in dadsops_rep, look at opus_rep
    if root_datasets[0][1] == "NULL":
        program_id = rootname[1:4]
        query1 = "SELECT DISTINCT proposal_id FROM executed WHERE "\
        "program_id='{0}'\ngo".format(program_id)
        SETTINGS["database"] = "opus_rep"
        prop = janky_connect(SETTINGS, query1)
        for i in range(len(root_datasets)):
            root_datasets[i][1] = prop

    return root_datasets

#-----------------------------------------------------------------------------#
#-----------------------------------------------------------------------------#

def parse_input(args):
    '''
    Parse the input argument to remove commas or * if present.

    Parameters:
    -----------
        args : list
            argparse list of input parameters

    Returns:
    --------
        ind_args : list
            List, each element is a cleaned up parameter.
    '''

    ind_args = []
    for item in args:
        if "," in item:
            tmp = (x for x in item.split(",") if x)
            for tmpitem in tmp:
                ind_args.append(tmpitem.strip("*"))
        else:
            ind_args.append(item.strip("*"))

    return ind_args

#-----------------------------------------------------------------------------#
#-----------------------------------------------------------------------------#

def is_proposal(mystring):
    '''
    Determine if an input is a proposal (e.g. 12046) or a rootname
    (e.g. LD201020)

    Parameters:
    -----------
        mystring : str
            String which describes data to be retrieved.

    Returns:
    --------
        status : Boolean
            True if mystring is a proposal ID.
    '''

    try:
        prop = int(mystring)
        status = True
    except ValueError:
        status = False

    return status

#-----------------------------------------------------------------------------#
#-----------------------------------------------------------------------------#
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # This is a required argument. To input multiple arguments, they must be
    # SPACE separated, not comma separated.
    parser.add_argument("data", nargs="+",  help="Reprocess programs or datasets")
    args = parser.parse_args()

    data = parse_input(args.data)
    print(data)

    all_data = []
    for item in data:
        isprop = is_proposal(item)
        if isprop:
            prop_datasets = query_proposal(item)
            all_data += prop_datasets
        else:
            root_datasets = handle_datasets(item)
            all_data += root_datasets

    to_retrieve = {row[0]:int(row[1]) for row in all_data}
    prop_keys = list(set(to_retrieve.values()))
    prop_vals = [[] for x in xrange(len(prop_keys))]
    prop_dict = dict(zip(prop_keys, prop_vals))
    for key in to_retrieve.keys():
        prop_dict[to_retrieve[key]].append(key)

    run_all_retrievals(prop_dict=prop_dict)
<commit_msg>Added CCI retrieval compatibility<commit_after>#! /usr/bin/env python
from __future__ import print_function, absolute_import, division
'''
Re-request and process COS data to keep products up to date.
'''

import argparse
import yaml
import os
import pdb

from .find_new_cos_data import janky_connect
from .request_data import run_all_retrievals

#-----------------------------------------------------------------------------#
#-----------------------------------------------------------------------------#
def query_proposal(proposal):
    '''
    Find all member datasets of a given proposal ID.

    Parameters:
    -----------
        proposal : int
            The proposal ID for which you wish to retrieve all data.

    Returns:
    --------
        prop_datasets : list
            List of all datasets given a proposal ID.
    '''

    # Open the configuration file for the MAST database connection (TSQL).
    config_file = os.path.join(os.environ['HOME'], "configure2.yaml")
    with open(config_file, 'r') as f:
        SETTINGS = yaml.load(f)

    # Need to use opus_rep DB instead of dadsops_rep to get ALL datasets
    SETTINGS["database"] = "opus_rep"
    query0 = "SELECT DISTINCT program_id FROM executed "\
    "WHERE proposal_id='{0}'\ngo".format(str(proposal))
    vid = janky_connect(SETTINGS, query0)

    # Now we need to go back to dadsop_rep to get ads_data_set_name
    SETTINGS["database"] = "dadsops_rep"
    query1 = "SELECT DISTINCT ads_data_set_name FROM "\
    "archive_data_set_all WHERE ads_data_set_name "\
    "LIKE '_{0}%'\ngo".format(str(vid[0]))
    datasets = janky_connect(SETTINGS, query1)
    prop_datasets = [[x, proposal] for x in datasets]

    return prop_datasets

#-----------------------------------------------------------------------------#
#-----------------------------------------------------------------------------#
def handle_datasets(rootname):
    '''
    Get proposal IDs and, if applicable, all member datasets given a visit or
    program.

    Parameters:
    -----------
        rootname : str
            The rootname (partial or full) to be retrieved.

    Returns:
    --------
        root_datasets : list
            List of all datasets for a given rootname.
    '''

    config_file = os.path.join(os.environ['HOME'], "configure2.yaml")
    with open(config_file, 'r') as f:
        SETTINGS = yaml.load(f)

    query0 = "SELECT DISTINCT ads_data_set_name,ads_pep_id FROM "\
    "archive_data_set_all WHERE ads_data_set_name "\
    "LIKE '{0}%'\ngo".format(str(rootname))
    root_datasets = janky_connect(SETTINGS, query0)

    # If data do not have a PID in dadsops_rep, look at opus_rep
    if root_datasets[0][1] == "NULL":
        if root_datasets[0][0].startswith("L_"):
            prop = "CCI"
        else:
            program_id = rootname[1:4]
            query1 = "SELECT DISTINCT proposal_id FROM executed WHERE "\
            "program_id='{0}'\ngo".format(program_id)
            SETTINGS["database"] = "opus_rep"
            prop = janky_connect(SETTINGS, query1)
            if not prop:
                prop = "NULL"
        for i in range(len(root_datasets)):
            root_datasets[i][1] = prop

    return root_datasets

#-----------------------------------------------------------------------------#
#-----------------------------------------------------------------------------#
def handle_cci(get_old=False):
    '''
    Retrieve all CCI datasets.

    Parameters:
    -----------
        get_old : Boolean
            In 2014, the CCI naming convention changed. Should the old CCIs
            be retrieved as well?

    Returns:
    --------
        cci_datasets : list
            List of all CCI datasets to be retrieved..
    '''

    config_file = os.path.join(os.environ['HOME'], "configure2.yaml")
    with open(config_file, 'r') as f:
        SETTINGS = yaml.load(f)

    query0 = "SELECT ads_data_set_name FROM archive_data_set_all WHERE "\
    "ads_archive_class='CSI'\ngo"
    all_ccis = janky_connect(SETTINGS, query0)

    if not get_old:
        cci_datasets = [[x,"CCI"] for x in all_ccis if x.startswith("L_")]
    else:
        cci_datasets = [[x,"CCI"] for x in all_ccis]

    return cci_datasets

#-----------------------------------------------------------------------------#
#-----------------------------------------------------------------------------#

def parse_input(args):
    '''
    Parse the input argument to remove commas or * if present.

    Parameters:
    -----------
        args : list
            argparse list of input parameters

    Returns:
    --------
        ind_args : list
            List, each element is a cleaned up parameter.
    '''

    ind_args = []
    for item in args:
        if "," in item:
            tmp = (x for x in item.split(",") if x)
            for tmpitem in tmp:
                ind_args.append(tmpitem.strip("*"))
        else:
            ind_args.append(item.strip("*"))

    return ind_args

#-----------------------------------------------------------------------------#
#-----------------------------------------------------------------------------#

def is_proposal(mystring):
    '''
    Determine if an input is a proposal (e.g. 12046) or a rootname
    (e.g. LD201020)

    Parameters:
    -----------
        mystring : str
            String which describes data to be retrieved.

    Returns:
    --------
        status : Boolean
            True if mystring is a proposal ID.
    '''

    try:
        prop = int(mystring)
        status = True
    except ValueError:
        status = False

    return status

#-----------------------------------------------------------------------------#
#-----------------------------------------------------------------------------#
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # This is a required argument. To input multiple arguments, they must be
    # SPACE separated, not comma separated.
    parser.add_argument("data", nargs="+",  help="Reprocess programs or datasets")
    args = parser.parse_args()

    data = parse_input(args.data)
    print("Requesting: {0}".format(data))

    all_data = []
    for item in data:
        if item.upper() == "CCI":
            all_data += handle_cci()
            continue
        isprop = is_proposal(item)
        if isprop:
            prop_datasets = query_proposal(item)
            all_data += prop_datasets
        else:
            root_datasets = handle_datasets(item)
            all_data += root_datasets

    to_retrieve = { (row[0]):(row[1] if row[1]=="CCI" else int(row[1])) for row in all_data}
    prop_keys = list(set(to_retrieve.values()))
    prop_vals = [[] for x in xrange(len(prop_keys))]
    prop_dict = dict(zip(prop_keys, prop_vals))
    for key in to_retrieve.keys():
        prop_dict[to_retrieve[key]].append(key)

    run_all_retrievals(prop_dict=prop_dict)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

""" Auto Encoder Example.
Using an auto encoder on MNIST handwritten digits.
References:
    Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner. "Gradient-based
    learning applied to document recognition." Proceedings of the IEEE,
    86(11):2278-2324, November 1998.
Links:
    [MNIST Dataset] http://yann.lecun.com/exdb/mnist/
"""
from __future__ import division, print_function, absolute_import

import argparse
import os
import sys
from datetime import datetime
import pickle
import json
import tempfile
import pprint
import yaml
import numpy as np
from hyperopt import hp, fmin, tpe, space_eval, Trials
from hyperopt import STATUS_OK, STATUS_FAIL
from hyperopt import pyll
import transform.logging_utils as lu
from transform.cfg_utils import load_config
import run_mlp as script
logger = None

if sys.version_info.major < 3:  # Python 2?
    # Using exec avoids a SyntaxError in Python 3.
    exec("""def reraise(exc_type, exc_value, exc_traceback=None):
                raise exc_type, exc_value, exc_traceback""")
else:
    def reraise(exc_type, exc_value, exc_traceback=None):
        if exc_value is None:
            exc_value = exc_type()
        if exc_value.__traceback__ is not exc_traceback:
            raise exc_value.with_traceback(exc_traceback)
        raise exc_value

def get_items_suffix(items):
    return '_'.join(['-'.join([k, str(v)]) for k, v in items])

def unique_dirname(dir_):
    dir_parent = os.path.dirname(dir_)
    prefix = os.path.basename(dir_) + '_'
    dir_ = tempfile.mkdtemp(suffix='', prefix=prefix, dir=dir_parent)
    logger.debug("Created directory %s", dir_)
    return dir_

def update_cfg(base, params):
    excl = ['base', 'run_dir', 'run_name', 'log_dir', 'pass_on_args']
    items_new = []
    for key, value in params.items():
        if key not in excl:
            base[key] = value
            items_new.append([key, value])
    for key, value in base.items():
        if isinstance(value, tuple):
            base[key] = list(value)
    suffix = params['run_name'] + '_' + get_items_suffix(items_new).replace(' ', '')
    run_dir_new = unique_dirname(os.path.join(params['run_dir'], suffix))
    base['log_dir'] = run_dir_new # keep log and run dirs the same
    base['run_dir'] = run_dir_new
    return base, items_new, suffix

def objective(params):
    status = STATUS_OK
    cfg, _, suffix = update_cfg(params['base'], params)
    fpath_cfg_dst = os.path.join(params['run_dir'], 'config_%s.yml' % suffix)
    logger.debug("Write config %s" % (fpath_cfg_dst))
    with open(fpath_cfg_dst, 'w') as h:
        h.write(yaml.dump(cfg))
    ch_args_in = ['-c', fpath_cfg_dst,
                  '--log_dir', cfg['log_dir'],
                  '--run_name', suffix,
                  '--run_dir', cfg['run_dir'],
                  ]
    for k, v in params['pass_on_args']:
        logger.debug("Passing on %s %s" % (k, str(v)))
        ch_args_in.extend([k, str(v)])
    args_ch = script.handleArgs(args=ch_args_in)
    result_run = \
        script.run(args_ch.run_name,
           args_ch
           )
    if result_run is None:
        status = STATUS_FAIL
    if status == STATUS_FAIL:
        logger.error("fmin failed with params: %s" % params)
    result = {
        "name"      : result_run.name,
        "loss"      : -result_run.max,
        "performance" : result_run.max,
        "status"    : status,
        "space"     : cfg,
    }
    return result

def add_space(space_base, space_dir):
    sys.path.append(args.space_dir)
    try:
        import space
        fpath_space = space.__file__
        from space import space as space_add
    except ImportError as e:
        module_name = str(e).split(' ')[-1]
        expected_module_path = os.path.join(args.space_dir, module_name + '.py')
        new_msg = str(e) + '. Verify that the module %s exists.' % expected_module_path
        logger.exception(new_msg)
        reraise(type(e), type(e)(
                new_msg), sys.exc_info()[2])
    logger.debug("Adding space definitions from %s" % fpath_space)
    for key, value in space_add.items():
        space_base[key] = value
    return space_base

def init_trials(fpath_trials, force_fresh=False):
    if not force_fresh:
        try:
            trials = pickle.load(open(fpath_trials, "rb"))
            logger.info("Loading saved trials from %s" % fpath_trials)
            logger.debug("Rerunning from {} trials to add more.".format(
                len(trials.trials)))
        except:
            trials = Trials()
            logger.info("Starting trials from scratch.")
    else:
        trials = Trials()
        logger.info("Starting trials from scratch.")
    return trials

def run(run_name, args):
    if args.run_dir is None:
        run_dir = os.path.join(args.log_dir, run_name)
    else:
        run_dir = args.run_dir
    run_dir_already_exists = False
    if not os.path.isdir(run_dir):
        os.makedirs(run_dir)
    else:
        run_dir_already_exists = True
    global logger
    logger = lu.setup_logging(os.path.join(run_dir, 'log.txt'),
                              name=[args.logger_name, None][args.logger_name_none])
    if run_dir_already_exists:
        logger.debug("Found run directory %s", run_dir)
    else:
        logger.debug("Created run directory %s", run_dir)
    fpath_trials = os.path.join(run_dir, "trials.pkl")
    trials = init_trials(fpath_trials, force_fresh=args.force_fresh_trials)
    max_evals = args.nb_evals
    try:
        trials = pickle.load(open(fpath_trials, "rb"))
        logger.info("Loading saved trials from %s" % fpath_trials)
        max_evals += len(trials.trials)
        logger.debug("Rerunning from {} trials to add more.".format(
            len(trials.trials)))
    except:
        trials = Trials()
        logger.info("Starting trials from scratch.")

    cfg = load_config(args.fpath_cfg, logger)
    space = {
        'log_dir' : os.path.abspath(os.path.expanduser(args.log_dir)),
        'run_dir' : run_dir,
        'run_name': run_name,
        'base'  : cfg,
        'pass_on_args' : [
            ['--per_process_gpu_memory_fraction', args.per_process_gpu_memory_fraction],
            ['--data_dir', args.data_dir],
            ],
        'tf_record_prefix'  : args.tf_record_prefix,
#        'learning_rate' : hp.choice('learning_rate', [0.5, 0.1, 0.05, 0.01, 0.005, 0.001, 0.0005, 0.0001, 0.00005, 0.00001, 0.000005, 0.000001]),
#        'lambda_l2' : hp.choice('lambda_l2', [0.0, 0.5, 0.1, 0.05, 0.01, 0.005, 0.001, 0.0005, 0.0001]),
        }
    space = add_space(space, args.space_dir)
    pp = pprint.PrettyPrinter(indent=4, width=100)
    for _ in range(10):
        pp.pprint(pyll.stochastic.sample(space))
    fpath_space = os.path.join(run_dir, "space.pkl")
    logger.debug("Save space to %s" % fpath_space)
    with open(fpath_space, "wb") as h:
        pickle.dump(space, h)
    best = fmin(fn=objective,
                space=space,
                algo=tpe.suggest,
                trials=trials,
                max_evals=max_evals)
    fpath_best_results = os.path.join(run_dir, 'best_results.txt.json')
    logger.info("Best results saved to %s" % fpath_best_results)
    best_eval = space_eval(space, best)
    logger.info("Best results obtained via %s" % best_eval)
    with open(fpath_best_results, "wb") as h:
        json.dump(best_eval, h)
    logger.info("Save %d trials to %s" % (len(trials), fpath_trials))
    with open(fpath_trials, "wb") as h:
        pickle.dump(trials, h)
    logger.info("Finished hyper run %s" % run_name)

def handleArgs(args=None):
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--config",
                        dest="fpath_cfg", type=str, required=True,
                        help="Path to master config file")
    parser.add_argument("--log_dir", dest="log_dir", type=str,
                        default=os.path.abspath(os.path.curdir),
                        help="Set parent log directory for all runs")
    parser.add_argument("--logger_name", dest="logger_name", type=str,
                        default=__name__,
                        help="Set name for process logging")
    parser.add_argument('--logger_name_none', action='store_true')
    parser.add_argument("--run_name", dest="run_name", type=str,
                        default=datetime.now().strftime("%Y-%m-%d_%H-%M-%S"),
                        help="Set name for run")
    parser.add_argument("--run_name_prefix", dest="run_name_prefix", type=str,
                        default='',
                        help="Set prefix run name")
    parser.add_argument("--nb_evals", dest="nb_evals", type=int, required=True,
                        help="Maximum number of evaluations")
    parser.add_argument("--space_dir", dest="space_dir", type=str,
                        required=True,
                        help="Path to directory with module for space definition")
    parser.add_argument("--run_dir", dest="run_dir", type=str, default=None,
                        help="Set run directory")
    parser.add_argument("--per_process_gpu_memory_fraction", dest="per_process_gpu_memory_fraction",
                        type=float, default=1.,
                        help="Tensorflow's gpu option per_process_gpu_memory_fraction")
    parser.add_argument("--data_dir", dest="data_dir", type=str,
                        required=True,
                        help="Path to data directory")
    parser.add_argument("--tf_record_prefix", dest="tf_record_prefix", type=str,
                        help="filename prefix for tf records files")
    parser.add_argument("--force_fresh_trials",  action='store_true',
                        dest="force_fresh_trials",
                        help="Force to start new set of trials and not resume existing ones.")

    return parser.parse_args(args=args)

if __name__ == '__main__':
    args = handleArgs()
    run(args.run_name_prefix + args.run_name,
        args
        )
    pass
<commit_msg>fix loading from trials file<commit_after># -*- coding: utf-8 -*-

""" Auto Encoder Example.
Using an auto encoder on MNIST handwritten digits.
References:
    Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner. "Gradient-based
    learning applied to document recognition." Proceedings of the IEEE,
    86(11):2278-2324, November 1998.
Links:
    [MNIST Dataset] http://yann.lecun.com/exdb/mnist/
"""
from __future__ import division, print_function, absolute_import

import argparse
import os
import sys
from datetime import datetime
import pickle
import json
import tempfile
import pprint
import yaml
import numpy as np
from hyperopt import hp, fmin, tpe, space_eval, Trials
from hyperopt import STATUS_OK, STATUS_FAIL
from hyperopt import pyll
import transform.logging_utils as lu
from transform.cfg_utils import load_config
import run_mlp as script
logger = None

if sys.version_info.major < 3:  # Python 2?
    # Using exec avoids a SyntaxError in Python 3.
    exec("""def reraise(exc_type, exc_value, exc_traceback=None):
                raise exc_type, exc_value, exc_traceback""")
else:
    def reraise(exc_type, exc_value, exc_traceback=None):
        if exc_value is None:
            exc_value = exc_type()
        if exc_value.__traceback__ is not exc_traceback:
            raise exc_value.with_traceback(exc_traceback)
        raise exc_value

def get_items_suffix(items):
    return '_'.join(['-'.join([k, str(v)]) for k, v in items])

def unique_dirname(dir_):
    dir_parent = os.path.dirname(dir_)
    prefix = os.path.basename(dir_) + '_'
    dir_ = tempfile.mkdtemp(suffix='', prefix=prefix, dir=dir_parent)
    logger.debug("Created directory %s", dir_)
    return dir_

def update_cfg(base, params):
    excl = ['base', 'run_dir', 'run_name', 'log_dir', 'pass_on_args']
    items_new = []
    for key, value in params.items():
        if key not in excl:
            base[key] = value
            items_new.append([key, value])
    for key, value in base.items():
        if isinstance(value, tuple):
            base[key] = list(value)
    suffix = params['run_name'] + '_' + get_items_suffix(items_new).replace(' ', '')
    run_dir_new = unique_dirname(os.path.join(params['run_dir'], suffix))
    base['log_dir'] = run_dir_new # keep log and run dirs the same
    base['run_dir'] = run_dir_new
    return base, items_new, suffix

def objective(params):
    status = STATUS_OK
    cfg, _, suffix = update_cfg(params['base'], params)
    fpath_cfg_dst = os.path.join(params['run_dir'], 'config_%s.yml' % suffix)
    logger.debug("Write config %s" % (fpath_cfg_dst))
    with open(fpath_cfg_dst, 'w') as h:
        h.write(yaml.dump(cfg))
    ch_args_in = ['-c', fpath_cfg_dst,
                  '--log_dir', cfg['log_dir'],
                  '--run_name', suffix,
                  '--run_dir', cfg['run_dir'],
                  ]
    for k, v in params['pass_on_args']:
        logger.debug("Passing on %s %s" % (k, str(v)))
        ch_args_in.extend([k, str(v)])
    args_ch = script.handleArgs(args=ch_args_in)
    result_run = \
        script.run(args_ch.run_name,
           args_ch
           )
    if result_run is None:
        status = STATUS_FAIL
    if status == STATUS_FAIL:
        logger.error("fmin failed with params: %s" % params)
    result = {
        "name"      : result_run.name,
        "loss"      : -result_run.max,
        "performance" : result_run.max,
        "status"    : status,
        "space"     : cfg,
    }
    return result

def add_space(space_base, space_dir):
    sys.path.append(args.space_dir)
    try:
        import space
        fpath_space = space.__file__
        from space import space as space_add
    except ImportError as e:
        module_name = str(e).split(' ')[-1]
        expected_module_path = os.path.join(args.space_dir, module_name + '.py')
        new_msg = str(e) + '. Verify that the module %s exists.' % expected_module_path
        logger.exception(new_msg)
        reraise(type(e), type(e)(
                new_msg), sys.exc_info()[2])
    logger.debug("Adding space definitions from %s" % fpath_space)
    for key, value in space_add.items():
        space_base[key] = value
    return space_base

def init_trials(fpath_trials, force_fresh=False):
    if not force_fresh:
        try:
            trials = pickle.load(open(fpath_trials, "rb"))
            logger.info("Loading saved trials from %s" % fpath_trials)
            logger.debug("Rerunning from {} trials to add more.".format(
                len(trials.trials)))
        except:
            trials = Trials()
            logger.info("Starting trials from scratch.")
    else:
        trials = Trials()
        logger.info("Starting trials from scratch.")
    return trials

def run(run_name, args):
    if args.run_dir is None:
        run_dir = os.path.join(args.log_dir, run_name)
    else:
        run_dir = args.run_dir
    run_dir_already_exists = False
    if not os.path.isdir(run_dir):
        os.makedirs(run_dir)
    else:
        run_dir_already_exists = True
    global logger
    logger = lu.setup_logging(os.path.join(run_dir, 'log.txt'),
                              name=[args.logger_name, None][args.logger_name_none])
    if run_dir_already_exists:
        logger.debug("Found run directory %s", run_dir)
    else:
        logger.debug("Created run directory %s", run_dir)
    fpath_trials = os.path.join(run_dir, "trials.pkl")
    trials = init_trials(fpath_trials, force_fresh=args.force_fresh_trials)
    max_evals = args.nb_evals + len(trials.trials)

    cfg = load_config(args.fpath_cfg, logger)
    space = {
        'log_dir' : os.path.abspath(os.path.expanduser(args.log_dir)),
        'run_dir' : run_dir,
        'run_name': run_name,
        'base'  : cfg,
        'pass_on_args' : [
            ['--per_process_gpu_memory_fraction', args.per_process_gpu_memory_fraction],
            ['--data_dir', args.data_dir],
            ],
        'tf_record_prefix'  : args.tf_record_prefix,
#        'learning_rate' : hp.choice('learning_rate', [0.5, 0.1, 0.05, 0.01, 0.005, 0.001, 0.0005, 0.0001, 0.00005, 0.00001, 0.000005, 0.000001]),
#        'lambda_l2' : hp.choice('lambda_l2', [0.0, 0.5, 0.1, 0.05, 0.01, 0.005, 0.001, 0.0005, 0.0001]),
        }
    space = add_space(space, args.space_dir)
    pp = pprint.PrettyPrinter(indent=4, width=100)
    for _ in range(10):
        pp.pprint(pyll.stochastic.sample(space))
    fpath_space = os.path.join(run_dir, "space.pkl")
    logger.debug("Save space to %s" % fpath_space)
    with open(fpath_space, "wb") as h:
        pickle.dump(space, h)
    best = fmin(fn=objective,
                space=space,
                algo=tpe.suggest,
                trials=trials,
                max_evals=max_evals)
    fpath_best_results = os.path.join(run_dir, 'best_results.txt.json')
    logger.info("Best results saved to %s" % fpath_best_results)
    best_eval = space_eval(space, best)
    logger.info("Best results obtained via %s" % best_eval)
    with open(fpath_best_results, "wb") as h:
        json.dump(best_eval, h)
    logger.info("Save %d trials to %s" % (len(trials), fpath_trials))
    with open(fpath_trials, "wb") as h:
        pickle.dump(trials, h)
    logger.info("Finished hyper run %s" % run_name)

def handleArgs(args=None):
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--config",
                        dest="fpath_cfg", type=str, required=True,
                        help="Path to master config file")
    parser.add_argument("--log_dir", dest="log_dir", type=str,
                        default=os.path.abspath(os.path.curdir),
                        help="Set parent log directory for all runs")
    parser.add_argument("--logger_name", dest="logger_name", type=str,
                        default=__name__,
                        help="Set name for process logging")
    parser.add_argument('--logger_name_none', action='store_true')
    parser.add_argument("--run_name", dest="run_name", type=str,
                        default=datetime.now().strftime("%Y-%m-%d_%H-%M-%S"),
                        help="Set name for run")
    parser.add_argument("--run_name_prefix", dest="run_name_prefix", type=str,
                        default='',
                        help="Set prefix run name")
    parser.add_argument("--nb_evals", dest="nb_evals", type=int, required=True,
                        help="Maximum number of evaluations")
    parser.add_argument("--space_dir", dest="space_dir", type=str,
                        required=True,
                        help="Path to directory with module for space definition")
    parser.add_argument("--run_dir", dest="run_dir", type=str, default=None,
                        help="Set run directory")
    parser.add_argument("--per_process_gpu_memory_fraction", dest="per_process_gpu_memory_fraction",
                        type=float, default=1.,
                        help="Tensorflow's gpu option per_process_gpu_memory_fraction")
    parser.add_argument("--data_dir", dest="data_dir", type=str,
                        required=True,
                        help="Path to data directory")
    parser.add_argument("--tf_record_prefix", dest="tf_record_prefix", type=str,
                        help="filename prefix for tf records files")
    parser.add_argument("--force_fresh_trials",  action='store_true',
                        dest="force_fresh_trials",
                        help="Force to start new set of trials and not resume existing ones.")

    return parser.parse_args(args=args)

if __name__ == '__main__':
    args = handleArgs()
    run(args.run_name_prefix + args.run_name,
        args
        )
    pass
<|endoftext|>


<commit_before><commit_msg>REF: fixed issue with file name with otsu<commit_after><|endoftext|>


<commit_before>#!/usr/bin/env python2
"""
Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
You may obtain a copy of the License at: [package root]/LICENSE.txt
Unless required by applicable law or agreed to in writing, software distributed under the License
 is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.

show-features

Show the features that capa extracts from the given sample,
to assist with the development of rules.

If you have a function with a capability that you'd like to detect,
you can run this tool and grep for the function/basic block/instruction addresses
to see what capa picks up.
This way, you can verify that capa successfully notices the features you'd reference.

Example::

    $ python scripts/show-features.py /tmp/suspicious.dll_
    ...
    file: 0x10004e4d: export(__entry)
    file: 0x10004706: export(Install)
    file: 0x10004c2b: export(uninstallA)
    file: 0x10005034: import(kernel32.GetStartupInfoA)
    file: 0x10005034: import(GetStartupInfoA)
    file: 0x10005048: import(kernel32.SetLastError)
    file: 0x00004e10: string(Y29ubmVjdA==)
    file: 0x00004e28: string(practicalmalwareanalysis.com)
    file: 0x00004e68: string(serve.html)
    file: 0x00004eb8: string(dW5zdXBwb3J0)
    file: 0x00004ec8: string(c2xlZXA=)
    func: 0x100012c2: characteristic(calls to)
    func: 0x10001000: characteristic(loop)
    bb  : 0x10001000: basic block
    insn: 0x10001000: mnemonic(push)
    insn: 0x10001001: mnemonic(push)
    insn: 0x10001002: mnemonic(push)
    insn: 0x10001003: mnemonic(push)
    insn: 0x10001004: mnemonic(push)
    insn: 0x10001005: mnemonic(push)
    insn: 0x10001006: mnemonic(xor)
    insn: 0x10001008: number(0x1)
    insn: 0x10001008: mnemonic(mov)
    bb  : 0x1000100a: basic block
    bb  : 0x1000100a: characteristic(tight loop)
    insn: 0x1000100a: mnemonic(movzx)
    insn: 0x1000100d: mnemonic(mov)
    insn: 0x1000100f: offset(0x1000A7C8)
    insn: 0x1000100f: mnemonic(mov)
    insn: 0x10001015: offset(0x100075C8)
    insn: 0x10001015: mnemonic(mov)
    insn: 0x1000101b: mnemonic(mov)
    insn: 0x1000101d: number(0x80)
    insn: 0x1000101d: mnemonic(and)
    insn: 0x10001020: mnemonic(neg)
    insn: 0x10001022: mnemonic(sbb)
    insn: 0x10001024: number(0x1B)
    insn: 0x10001024: mnemonic(and)
    insn: 0x10001027: number(0x1)
    insn: 0x10001027: mnemonic(shl)
    ...
"""
import sys
import logging
import argparse

import capa.main
import capa.rules
import capa.engine
import capa.helpers
import capa.features
import capa.features.freeze

logger = logging.getLogger("capa.show-features")

def main(argv=None):
    if argv is None:
        argv = sys.argv[1:]

    parser = argparse.ArgumentParser(description="Show the features that capa extracts from the given sample")
    capa.main.install_common_args(parser, wanted={"format", "sample"})

    parser.add_argument("-F", "--function", type=lambda x: int(x, 0x10), help="Show features for specific function")
    args = parser.parse_args(args=argv)
    capa.main.handle_common_args(args)

    try:
        taste = capa.helpers.get_file_taste(args.sample)
    except IOError as e:
        logger.error("%s", str(e))
        return -1

    if (args.format == "freeze") or (args.format == "auto" and capa.features.freeze.is_freeze(taste)):
        with open(args.sample, "rb") as f:
            extractor = capa.features.freeze.load(f.read())
    else:
        try:
            extractor = capa.main.get_extractor(args.sample, args.format, capa.main.BACKEND_VIV, sigpaths=[])
        except capa.main.UnsupportedFormatError:
            logger.error("-" * 80)
            logger.error(" Input file does not appear to be a PE file.")
            logger.error(" ")
            logger.error(
                " capa currently only supports analyzing PE files (or shellcode, when using --format sc32|sc64)."
            )
            logger.error(" If you don't know the input file type, you can try using the `file` utility to guess it.")
            logger.error("-" * 80)
            return -1
        except capa.main.UnsupportedRuntimeError:
            logger.error("-" * 80)
            logger.error(" Unsupported runtime or Python interpreter.")
            logger.error(" ")
            logger.error(" capa supports running under Python 2.7 using Vivisect for binary analysis.")
            logger.error(" It can also run within IDA Pro, using either Python 2.7 or 3.5+.")
            logger.error(" ")
            logger.error(" If you're seeing this message on the command line, please ensure you're running Python 2.7.")
            logger.error("-" * 80)
            return -1

    if not args.function:
        for feature, va in extractor.extract_file_features():
            if va:
                print("file: 0x%08x: %s" % (va, feature))
            else:
                print("file: 0x00000000: %s" % (feature))

    functions = extractor.get_functions()

    if args.function:
        if args.format == "freeze":
            functions = tuple(filter(lambda f: f == args.function, functions))
        else:
            functions = tuple(filter(lambda f: int(f) == args.function, functions))

            if args.function not in [int(f) for f in functions]:
                print("0x%X not a function" % args.function)
                return -1

        if len(functions) == 0:
            print("0x%X not a function")
            return -1

    print_features(functions, extractor)

    return 0

def ida_main():
    function = idc.get_func_attr(idc.here(), idc.FUNCATTR_START)
    print("getting features for current function 0x%X" % function)

    extractor = capa.features.extractors.ida.IdaFeatureExtractor()

    if not function:
        for feature, va in extractor.extract_file_features():
            if va:
                print("file: 0x%08x: %s" % (va, feature))
            else:
                print("file: 0x00000000: %s" % (feature))
        return

    functions = extractor.get_functions()

    if function:
        functions = tuple(filter(lambda f: f.start_ea == function, functions))

        if len(functions) == 0:
            print("0x%X not a function" % function)
            return -1

    print_features(functions, extractor)

    return 0

def print_features(functions, extractor):
    for f in functions:
        for feature, va in extractor.extract_function_features(f):
            print("func: 0x%08x: %s" % (va, feature))

        for bb in extractor.get_basic_blocks(f):
            for feature, va in extractor.extract_basic_block_features(f, bb):
                print("bb  : 0x%08x: %s" % (va, feature))

            for insn in extractor.get_instructions(f, bb):
                for feature, va in extractor.extract_insn_features(f, bb, insn):
                    try:
                        print("insn: 0x%08x: %s" % (va, feature))
                    except UnicodeEncodeError:
                        # may be an issue while piping to less and encountering non-ascii characters
                        continue

if __name__ == "__main__":
    if capa.main.is_runtime_ida():
        ida_main()
    else:
        sys.exit(main())
<commit_msg>show-features: accept signatures or use default<commit_after>#!/usr/bin/env python2
"""
Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
You may obtain a copy of the License at: [package root]/LICENSE.txt
Unless required by applicable law or agreed to in writing, software distributed under the License
 is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.

show-features

Show the features that capa extracts from the given sample,
to assist with the development of rules.

If you have a function with a capability that you'd like to detect,
you can run this tool and grep for the function/basic block/instruction addresses
to see what capa picks up.
This way, you can verify that capa successfully notices the features you'd reference.

Example::

    $ python scripts/show-features.py /tmp/suspicious.dll_
    ...
    file: 0x10004e4d: export(__entry)
    file: 0x10004706: export(Install)
    file: 0x10004c2b: export(uninstallA)
    file: 0x10005034: import(kernel32.GetStartupInfoA)
    file: 0x10005034: import(GetStartupInfoA)
    file: 0x10005048: import(kernel32.SetLastError)
    file: 0x00004e10: string(Y29ubmVjdA==)
    file: 0x00004e28: string(practicalmalwareanalysis.com)
    file: 0x00004e68: string(serve.html)
    file: 0x00004eb8: string(dW5zdXBwb3J0)
    file: 0x00004ec8: string(c2xlZXA=)
    func: 0x100012c2: characteristic(calls to)
    func: 0x10001000: characteristic(loop)
    bb  : 0x10001000: basic block
    insn: 0x10001000: mnemonic(push)
    insn: 0x10001001: mnemonic(push)
    insn: 0x10001002: mnemonic(push)
    insn: 0x10001003: mnemonic(push)
    insn: 0x10001004: mnemonic(push)
    insn: 0x10001005: mnemonic(push)
    insn: 0x10001006: mnemonic(xor)
    insn: 0x10001008: number(0x1)
    insn: 0x10001008: mnemonic(mov)
    bb  : 0x1000100a: basic block
    bb  : 0x1000100a: characteristic(tight loop)
    insn: 0x1000100a: mnemonic(movzx)
    insn: 0x1000100d: mnemonic(mov)
    insn: 0x1000100f: offset(0x1000A7C8)
    insn: 0x1000100f: mnemonic(mov)
    insn: 0x10001015: offset(0x100075C8)
    insn: 0x10001015: mnemonic(mov)
    insn: 0x1000101b: mnemonic(mov)
    insn: 0x1000101d: number(0x80)
    insn: 0x1000101d: mnemonic(and)
    insn: 0x10001020: mnemonic(neg)
    insn: 0x10001022: mnemonic(sbb)
    insn: 0x10001024: number(0x1B)
    insn: 0x10001024: mnemonic(and)
    insn: 0x10001027: number(0x1)
    insn: 0x10001027: mnemonic(shl)
    ...
"""
import sys
import logging
import argparse

import capa.main
import capa.rules
import capa.engine
import capa.helpers
import capa.features
import capa.features.freeze

logger = logging.getLogger("capa.show-features")

def main(argv=None):
    if argv is None:
        argv = sys.argv[1:]

    parser = argparse.ArgumentParser(description="Show the features that capa extracts from the given sample")
    capa.main.install_common_args(parser, wanted={"format", "sample", "signatures"})

    parser.add_argument("-F", "--function", type=lambda x: int(x, 0x10), help="Show features for specific function")
    args = parser.parse_args(args=argv)
    capa.main.handle_common_args(args)

    try:
        taste = capa.helpers.get_file_taste(args.sample)
    except IOError as e:
        logger.error("%s", str(e))
        return -1

    if (args.format == "freeze") or (args.format == "auto" and capa.features.freeze.is_freeze(taste)):
        with open(args.sample, "rb") as f:
            extractor = capa.features.freeze.load(f.read())
    else:
        try:
            extractor = capa.main.get_extractor(args.sample, args.format, capa.main.BACKEND_VIV, sigpaths=args.signatures)
        except capa.main.UnsupportedFormatError:
            logger.error("-" * 80)
            logger.error(" Input file does not appear to be a PE file.")
            logger.error(" ")
            logger.error(
                " capa currently only supports analyzing PE files (or shellcode, when using --format sc32|sc64)."
            )
            logger.error(" If you don't know the input file type, you can try using the `file` utility to guess it.")
            logger.error("-" * 80)
            return -1
        except capa.main.UnsupportedRuntimeError:
            logger.error("-" * 80)
            logger.error(" Unsupported runtime or Python interpreter.")
            logger.error(" ")
            logger.error(" capa supports running under Python 2.7 using Vivisect for binary analysis.")
            logger.error(" It can also run within IDA Pro, using either Python 2.7 or 3.5+.")
            logger.error(" ")
            logger.error(" If you're seeing this message on the command line, please ensure you're running Python 2.7.")
            logger.error("-" * 80)
            return -1

    if not args.function:
        for feature, va in extractor.extract_file_features():
            if va:
                print("file: 0x%08x: %s" % (va, feature))
            else:
                print("file: 0x00000000: %s" % (feature))

    functions = extractor.get_functions()

    if args.function:
        if args.format == "freeze":
            functions = tuple(filter(lambda f: f == args.function, functions))
        else:
            functions = tuple(filter(lambda f: int(f) == args.function, functions))

            if args.function not in [int(f) for f in functions]:
                print("0x%X not a function" % args.function)
                return -1

        if len(functions) == 0:
            print("0x%X not a function")
            return -1

    print_features(functions, extractor)

    return 0

def ida_main():
    function = idc.get_func_attr(idc.here(), idc.FUNCATTR_START)
    print("getting features for current function 0x%X" % function)

    extractor = capa.features.extractors.ida.IdaFeatureExtractor()

    if not function:
        for feature, va in extractor.extract_file_features():
            if va:
                print("file: 0x%08x: %s" % (va, feature))
            else:
                print("file: 0x00000000: %s" % (feature))
        return

    functions = extractor.get_functions()

    if function:
        functions = tuple(filter(lambda f: f.start_ea == function, functions))

        if len(functions) == 0:
            print("0x%X not a function" % function)
            return -1

    print_features(functions, extractor)

    return 0

def print_features(functions, extractor):
    for f in functions:
        for feature, va in extractor.extract_function_features(f):
            print("func: 0x%08x: %s" % (va, feature))

        for bb in extractor.get_basic_blocks(f):
            for feature, va in extractor.extract_basic_block_features(f, bb):
                print("bb  : 0x%08x: %s" % (va, feature))

            for insn in extractor.get_instructions(f, bb):
                for feature, va in extractor.extract_insn_features(f, bb, insn):
                    try:
                        print("insn: 0x%08x: %s" % (va, feature))
                    except UnicodeEncodeError:
                        # may be an issue while piping to less and encountering non-ascii characters
                        continue

if __name__ == "__main__":
    if capa.main.is_runtime_ida():
        ida_main()
    else:
        sys.exit(main())
<|endoftext|>


<commit_before>"""
Purpose: to compute the chi_heap recursion equation of Harte 2007
and to export the results as a .csv file
"""

import mete
import numpy as np
import csv
import sys

if len(sys.argv) > 1:
    A = int(sys.argv[1])
    A0 = int(sys.argv[2])
    shape = sys.argv[3]
    sadType = sys.argv[4]
    abu_filepath = sys.argv[5]
    out_filepath = sys.argv[6]
    if len(sys.argv) == 8:
        unit_distance = float(sys.argv[7])
    else:
        unit_distance = 1
else:
    print 'Error: Specify A, A0, shape, and abu file at the command line'

if abu_filepath  != 'None':
    datafile = open(abu_filepath, 'r')
    datareader = csv.reader(datafile)
    data = []
    for row in datareader:
        data.append(row)
    abu = [int(x) for x in data[0]]
else:
    print 'Error: Abundance data not specified'

if sadType == 'meteSAD':
    S0 = len(abu)
    N0 = sum(abu)
    sor = mete.sor_heap(A, A0, S0, N0, shape, unit_distance)
else:
    sor = mete.sor_heap_fixed_abu(A, abu, A0, shape, unit_distance)

writer = open(out_filepath, 'wb')
datawriter = csv.writer(writer)

datawriter.writerow(['j', 'dist', 'sor'])
for i in range(0, np.shape(sor)[0]):
    datawriter.writerow(sor[i, ])

writer.close()
<commit_msg>Revised so that multiple spatial grains can be specified from the command line.<commit_after>"""
Purpose: to compute the chi_heap recursion equation of Harte 2007
and to export the results as a .csv file
"""

import mete
import numpy as np
import csv
import sys

if len(sys.argv) > 1:
    A_vals = map(int, sys.argv[1])
    A0 = int(sys.argv[2])
    shape = sys.argv[3]
    sadType = sys.argv[4]
    abu_filepath = sys.argv[5]
    out_filepath = sys.argv[6]
    if len(sys.argv) == 8:
        unit_distance = float(sys.argv[7])
    else:
        unit_distance = 1
else:
    print 'Error: Specify A, A0, shape, and abu file at the command line'

if abu_filepath  != 'None':
    datafile = open(abu_filepath, 'r')
    datareader = csv.reader(datafile)
    data = []
    for row in datareader:
        data.append(row)
    abu = [int(x) for x in data[0]]
else:
    print 'Error: Abundance data not specified'

if sadType == 'meteSAD':
    S0 = len(abu)
    N0 = sum(abu)
    sor = [mete.sor_heap(A, A0, S0, N0, shape, unit_distance) for A in A_vals]
    sor = np.concatenate((sor), axis=0)
else:
    sor = [mete.sor_heap_fixed_abu(A, abu, A0, shape, unit_distance) for A in A_vals]
    sor = np.concatenate((sor), axis=0)

writer = open(out_filepath, 'wb')
datawriter = csv.writer(writer)

datawriter.writerow(['i', 'j', 'dist', 'sor'])
for i in range(0, np.shape(sor)[0]):
    datawriter.writerow(sor[i, ])

writer.close()

<|endoftext|>


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

import os
import sys
import socket
import urllib
import string
import random

class Signal:
    def __init__(self):
        self.slots = {}
    def ensure(self, func, context=""):
        if context not in self.slots:
            self.connect(func, context)
            return
        if func not in self.slots[context]:
            self.connect(func, context)
    def connect(self, func, context=""):
        if context not in self.slots:
            self.slots[context] = []
        self.slots[context] += [func]
    def clear(self):
        self.slots = {}
    def disconnect(self, context):
        del self.slots[context]
    def emit(self, *args, **kwargs):
        limit_context = kwargs.get("limit_context", None)
        brk = kwargs.get("allow_break", False)
        force_brk = kwargs.get("force_break", False)
        include_context = kwargs.get("include_context", False)
        for ctx, funcs in self.slots.items():
            if not limit_context or ctx in limit_context:
                for func in funcs:
                    r = None
                    if include_context:
                        r = func(ctx, *args)
                    else:
                        r = func(*args)
                    if brk and r:
                        return
                    if force_brk:
                        return
                    continue

class Server:
    def __init__(self, sock):
        self.sock = sock
        self.on_msg = Signal()
        self.on_data = Signal()
        self.on_command = Signal()
        self.on_enter = Signal()
        self.on_quit = Signal()
        self.test = not sock

    def send(self, msg):
        if self.test:
            print msg
        else:
            self.sock.send(msg)
    def say(self, dest, msg):
        if type(dest) == type([]):
            if TEST:
                print "(all) %s" % msg
            else:
                for chan in dest:
                    self.sock.send("PRIVMSG %s :%s\n" % (chan, msg))
            return
        if TEST:
            print "(%s) %s" % (dest,msg)
        else:
            self.sock.send("PRIVMSG %s :%s\n" % (dest, msg))
    def broadcast(self, msg):
        self.say(CHANS, msg)

config_path = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    "config.py"
)
with open(config_path) as source:
    eval(compile(source.read(), "config.py", 'exec'))

buf = ""

TEST = set(["-t","--test"]) & set(sys.argv[1:])

sock = None
if not TEST:
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
serv = Server(sock)
if not TEST:
    sock.connect((HOST, PORT))
    sock.send("NICK %s\n" % NICK)
    sock.send("USER %s %s %s :%s\n" % (IDENT, NICK, HOST, REALNAME))

plugins = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    "plugins"
)
for cmd in PLUGINS:
    print "loading %s plugin..." % cmd
    with open(os.path.join(plugins,cmd+".py")) as source:
        eval(compile(source.read(), "%s.py" % cmd, 'exec'))

#try:
#    with open("plugins/__init__.py") as source:
#        eval(compile(source.read(), "plugins.py", 'exec'))
#except:
#    sys.exit(1)

def about(cmd, serv, nick, dest, msg):
    serv.send("PRIVMSG %s :I am litebot (github.com/flipcoder/litebot)\n" % dest)
    serv.send("PRIVMSG %s :Commands (prefix with %%): %s\n" % (dest, ", ".join(sorted(serv.on_command.slots))))

logged_in = False or TEST

print "%s running" % NICK

if TEST:
    print 'Test mode on %s.' % CHANS[0]

try:
    while True:

        if not TEST:

            buf = sock.recv(4096)

            if buf.find("PING") != -1:
                sock.send("PONG %s\r\n" % buf.split()[1]+"\n")
                if not logged_in:
                    sock.send("PRIVMSG nickserv identify %s\n" % PASS)
                    for chan in CHANS:
                        sock.send("JOIN %s\n" % chan)
                    logged_in = True
                    print "signed in"
                    serv.on_enter.emit(serv,
                        include_context=True, allow_break=True)
                continue

            serv.on_data.emit(serv, buf, include_context=True, allow_break=True)

        if buf.find("PRIVMSG") != -1 or TEST:
            if not TEST:
                buf = buf.replace("!",":").split(":")
                nick = buf[1]
                if len(buf[3]) <= 2:
                    continue
                msg = ":".join(buf[3:])[:-2]
                dest = buf[2].split()[2]

            if TEST:
                nick = 'user'
                dest = CHANS[0]
                msg = raw_input('> ')

            if not TEST or msg:
                serv.on_msg.emit(serv, nick, dest, msg,
                    include_context=True, allow_break=True)

            if msg=="%" or msg=="%help":
                about("about", serv, nick, dest, msg)
                continue

            if msg and msg.startswith("%"):
                msg = msg[1:]
                msg = msg.strip()
                chop = msg.find(" ")
                if chop != -1:
                    cmd = msg[:chop]
                    msg = msg[chop+1:]
                else:
                    cmd = msg
                    msg = ""
                serv.on_command.emit(serv, nick, dest, msg,
                    include_context=True, limit_context=[cmd], force_break=True
                )
                continue
except:
    serv.on_quit.emit(serv, include_context=True)
    raise

<commit_msg>bool convert<commit_after>#!/usr/bin/env python2

import os
import sys
import socket
import urllib
import string
import random

class Signal:
    def __init__(self):
        self.slots = {}
    def ensure(self, func, context=""):
        if context not in self.slots:
            self.connect(func, context)
            return
        if func not in self.slots[context]:
            self.connect(func, context)
    def connect(self, func, context=""):
        if context not in self.slots:
            self.slots[context] = []
        self.slots[context] += [func]
    def clear(self):
        self.slots = {}
    def disconnect(self, context):
        del self.slots[context]
    def emit(self, *args, **kwargs):
        limit_context = kwargs.get("limit_context", None)
        brk = kwargs.get("allow_break", False)
        force_brk = kwargs.get("force_break", False)
        include_context = kwargs.get("include_context", False)
        for ctx, funcs in self.slots.items():
            if not limit_context or ctx in limit_context:
                for func in funcs:
                    r = None
                    if include_context:
                        r = func(ctx, *args)
                    else:
                        r = func(*args)
                    if brk and r:
                        return
                    if force_brk:
                        return
                    continue

class Server:
    def __init__(self, sock):
        self.sock = sock
        self.on_msg = Signal()
        self.on_data = Signal()
        self.on_command = Signal()
        self.on_enter = Signal()
        self.on_quit = Signal()
        self.test = not sock

    def send(self, msg):
        if self.test:
            print msg
        else:
            self.sock.send(msg)
    def say(self, dest, msg):
        if type(dest) == type([]):
            if TEST:
                print "(all) %s" % msg
            else:
                for chan in dest:
                    self.sock.send("PRIVMSG %s :%s\n" % (chan, msg))
            return
        if TEST:
            print "(%s) %s" % (dest,msg)
        else:
            self.sock.send("PRIVMSG %s :%s\n" % (dest, msg))
    def broadcast(self, msg):
        self.say(CHANS, msg)

config_path = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    "config.py"
)
with open(config_path) as source:
    eval(compile(source.read(), "config.py", 'exec'))

buf = ""

TEST = bool(set(["-t","--test"]) & set(sys.argv[1:]))

sock = None
if not TEST:
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
serv = Server(sock)
if not TEST:
    sock.connect((HOST, PORT))
    sock.send("NICK %s\n" % NICK)
    sock.send("USER %s %s %s :%s\n" % (IDENT, NICK, HOST, REALNAME))

plugins = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    "plugins"
)
for cmd in PLUGINS:
    print "loading %s plugin..." % cmd
    with open(os.path.join(plugins,cmd+".py")) as source:
        eval(compile(source.read(), "%s.py" % cmd, 'exec'))

#try:
#    with open("plugins/__init__.py") as source:
#        eval(compile(source.read(), "plugins.py", 'exec'))
#except:
#    sys.exit(1)

def about(cmd, serv, nick, dest, msg):
    serv.send("PRIVMSG %s :I am litebot (github.com/flipcoder/litebot)\n" % dest)
    serv.send("PRIVMSG %s :Commands (prefix with %%): %s\n" % (dest, ", ".join(sorted(serv.on_command.slots))))

logged_in = False or TEST

print "%s running" % NICK

if TEST:
    print 'Test mode on %s.' % CHANS[0]

try:
    while True:

        if not TEST:

            buf = sock.recv(4096)

            if buf.find("PING") != -1:
                sock.send("PONG %s\r\n" % buf.split()[1]+"\n")
                if not logged_in:
                    sock.send("PRIVMSG nickserv identify %s\n" % PASS)
                    for chan in CHANS:
                        sock.send("JOIN %s\n" % chan)
                    logged_in = True
                    print "signed in"
                    serv.on_enter.emit(serv,
                        include_context=True, allow_break=True)
                continue

            serv.on_data.emit(serv, buf, include_context=True, allow_break=True)

        if buf.find("PRIVMSG") != -1 or TEST:
            if not TEST:
                buf = buf.replace("!",":").split(":")
                nick = buf[1]
                if len(buf[3]) <= 2:
                    continue
                msg = ":".join(buf[3:])[:-2]
                dest = buf[2].split()[2]

            if TEST:
                nick = 'user'
                dest = CHANS[0]
                msg = raw_input('> ')

            if not TEST or msg:
                serv.on_msg.emit(serv, nick, dest, msg,
                    include_context=True, allow_break=True)

            if msg=="%" or msg=="%help":
                about("about", serv, nick, dest, msg)
                continue

            if msg and msg.startswith("%"):
                msg = msg[1:]
                msg = msg.strip()
                chop = msg.find(" ")
                if chop != -1:
                    cmd = msg[:chop]
                    msg = msg[chop+1:]
                else:
                    cmd = msg
                    msg = ""
                serv.on_command.emit(serv, nick, dest, msg,
                    include_context=True, limit_context=[cmd], force_break=True
                )
                continue
except:
    serv.on_quit.emit(serv, include_context=True)
    raise

<|endoftext|>