<commit_before>import logging

import pytest

import sdk_cmd as cmd
import sdk_install
import sdk_metrics
import sdk_upgrade
from tests.config import *

log = logging.getLogger(__name__)

FOLDERED_SERVICE_NAME = sdk_utils.get_foldered_name(PACKAGE_NAME)

@pytest.fixture(scope='module', autouse=True)
def configure_package(configure_universe):
    try:
        log.info("Ensure elasticsearch and kibana are uninstalled...")
        sdk_install.uninstall(KIBANA_PACKAGE_NAME)
        sdk_install.uninstall(FOLDERED_SERVICE_NAME, package_name=PACKAGE_NAME)
        sdk_utils.gc_frameworks()

        sdk_upgrade.test_upgrade(
            "beta-{}".format(PACKAGE_NAME),
            PACKAGE_NAME,
            DEFAULT_TASK_COUNT,
            service_name=FOLDERED_SERVICE_NAME,
            additional_options={"service": {"name": FOLDERED_SERVICE_NAME}})

        yield  # let the test session execute
    finally:
        log.info("Clean up elasticsearch and kibana...")
        sdk_install.uninstall(KIBANA_PACKAGE_NAME)
        sdk_install.uninstall(FOLDERED_SERVICE_NAME, package_name=PACKAGE_NAME)

@pytest.fixture(autouse=True)
def pre_test_setup():
    sdk_tasks.check_running(FOLDERED_SERVICE_NAME, DEFAULT_TASK_COUNT)
    wait_for_expected_nodes_to_exist(service_name=FOLDERED_SERVICE_NAME)

@pytest.fixture
def default_populated_index():
    delete_index(DEFAULT_INDEX_NAME, service_name=FOLDERED_SERVICE_NAME)
    create_index(DEFAULT_INDEX_NAME, DEFAULT_SETTINGS_MAPPINGS, service_name=FOLDERED_SERVICE_NAME)
    create_document(DEFAULT_INDEX_NAME, DEFAULT_INDEX_TYPE, 1, {"name": "Loren", "role": "developer"}, service_name=FOLDERED_SERVICE_NAME)

@pytest.mark.sanity
@pytest.mark.smoke
def test_service_health():
    assert shakedown.service_healthy(FOLDERED_SERVICE_NAME)

@pytest.mark.sanity
def test_endpoints():
    # check that we can reach the scheduler via admin router, and that returned endpoints are sanitized:
    for endpoint in ENDPOINT_TYPES:
        endpoints = json.loads(cmd.run_cli('elastic --name={} endpoints {}'.format(FOLDERED_SERVICE_NAME, endpoint)))
        host = endpoint.split('-')[0] # 'coordinator-http' => 'coordinator'
        assert endpoints['dns'][0].startswith(sdk_hosts.autoip_host(FOLDERED_SERVICE_NAME, host + '-0-node'))
        assert endpoints['vip'].startswith(sdk_hosts.vip_host(FOLDERED_SERVICE_NAME, host))

@pytest.mark.sanity
def test_indexing(default_populated_index):
    indices_stats = get_elasticsearch_indices_stats(DEFAULT_INDEX_NAME, service_name=FOLDERED_SERVICE_NAME)
    assert indices_stats["_all"]["primaries"]["docs"]["count"] == 1
    doc = get_document(DEFAULT_INDEX_NAME, DEFAULT_INDEX_TYPE, 1, service_name=FOLDERED_SERVICE_NAME)
    assert doc["_source"]["name"] == "Loren"

@pytest.mark.sanity
@pytest.mark.metrics
@sdk_utils.dcos_1_9_or_higher
def test_metrics():
    sdk_metrics.wait_for_any_metrics(FOLDERED_SERVICE_NAME, "data-0-node", DEFAULT_ELASTIC_TIMEOUT)

@pytest.mark.focus
@pytest.mark.sanity
@pytest.mark.timeout(60 * 60)
def test_xpack_toggle_with_kibana(default_populated_index):
    log.info("\n***** Verify X-Pack disabled by default in elasticsearch")
    verify_commercial_api_status(False, service_name=FOLDERED_SERVICE_NAME)

    log.info("\n***** Test kibana with X-Pack disabled...")
    shakedown.install_package(KIBANA_PACKAGE_NAME, options_json={
        "kibana": {
            "elasticsearch_url": "http://" + sdk_hosts.vip_host(FOLDERED_SERVICE_NAME, "coordinator", 9200)
        }})
    shakedown.deployment_wait(app_id="/{}".format(KIBANA_PACKAGE_NAME), timeout=DEFAULT_KIBANA_TIMEOUT)
    check_kibana_adminrouter_integration("service/{}/".format(KIBANA_PACKAGE_NAME))
    log.info("Uninstall kibana with X-Pack disabled")
    sdk_install.uninstall(KIBANA_PACKAGE_NAME)

    log.info("\n***** Set/verify X-Pack enabled in elasticsearch")
    enable_xpack(service_name=FOLDERED_SERVICE_NAME)
    verify_commercial_api_status(True, service_name=FOLDERED_SERVICE_NAME)
    verify_xpack_license(service_name=FOLDERED_SERVICE_NAME)

    log.info("\n***** Write some data while enabled, disable X-Pack, and verify we can still read what we wrote.")
    create_document(
        DEFAULT_INDEX_NAME,
        DEFAULT_INDEX_TYPE,
        2,
        {"name": "X-Pack", "role": "commercial plugin"},
        service_name=FOLDERED_SERVICE_NAME)

    log.info("\n***** Test kibana with X-Pack enabled...")
    shakedown.install_package(KIBANA_PACKAGE_NAME, options_json={
        "kibana": {
            "elasticsearch_url": "http://" + sdk_hosts.vip_host(FOLDERED_SERVICE_NAME, "coordinator", 9200),
            "xpack_enabled": True
        }})
    log.info("\n***** Installing Kibana w/X-Pack can take as much as 15 minutes for Marathon deployment ")
    log.info("to complete due to a configured HTTP health check. (typical: 12 minutes)")
    shakedown.deployment_wait(app_id="/{}".format(KIBANA_PACKAGE_NAME), timeout=DEFAULT_KIBANA_TIMEOUT)
    check_kibana_adminrouter_integration("service/{}/login".format(KIBANA_PACKAGE_NAME))
    log.info("\n***** Uninstall kibana with X-Pack enabled")
    sdk_install.uninstall(KIBANA_PACKAGE_NAME)

    log.info("\n***** Disable X-Pack in elasticsearch.")
    disable_xpack(service_name=FOLDERED_SERVICE_NAME)
    log.info("\n***** Verify we can still read what we wrote when X-Pack was enabled.")
    verify_commercial_api_status(False, service_name=FOLDERED_SERVICE_NAME)
    doc = get_document(DEFAULT_INDEX_NAME, DEFAULT_INDEX_TYPE, 2, service_name=FOLDERED_SERVICE_NAME)
    assert doc["_source"]["name"] == "X-Pack"

@pytest.mark.recovery
@pytest.mark.sanity
def test_losing_and_regaining_index_health(default_populated_index):
    check_elasticsearch_index_health(DEFAULT_INDEX_NAME, "green", service_name=FOLDERED_SERVICE_NAME)
    shakedown.kill_process_on_host(sdk_hosts.system_host(FOLDERED_SERVICE_NAME, "data-0-node"), "data__.*Elasticsearch")
    check_elasticsearch_index_health(DEFAULT_INDEX_NAME, "yellow", service_name=FOLDERED_SERVICE_NAME)
    check_elasticsearch_index_health(DEFAULT_INDEX_NAME, "green", service_name=FOLDERED_SERVICE_NAME)

@pytest.mark.recovery
@pytest.mark.sanity
def test_master_reelection():
    initial_master = get_elasticsearch_master(service_name=FOLDERED_SERVICE_NAME)
    shakedown.kill_process_on_host(sdk_hosts.system_host(FOLDERED_SERVICE_NAME, initial_master), "master__.*Elasticsearch")
    wait_for_expected_nodes_to_exist(service_name=FOLDERED_SERVICE_NAME)
    new_master = get_elasticsearch_master(service_name=FOLDERED_SERVICE_NAME)
    assert new_master.startswith("master") and new_master != initial_master

@pytest.mark.recovery
@pytest.mark.sanity
def test_master_node_replace():
    # Ideally, the pod will get placed on a different agent. This test will verify that the remaining two masters
    # find the replaced master at its new IP address. This requires a reasonably low TTL for Java DNS lookups.
    cmd.run_cli('elastic --name={} pod replace master-0'.format(FOLDERED_SERVICE_NAME))
    # setup_function will verify that the cluster becomes healthy again.

@pytest.mark.recovery
@pytest.mark.sanity
def test_plugin_install_and_uninstall(default_populated_index):
    plugin_name = 'analysis-phonetic'
    config = sdk_marathon.get_config(FOLDERED_SERVICE_NAME)
    config['env']['TASKCFG_ALL_ELASTICSEARCH_PLUGINS'] = plugin_name
    sdk_marathon.update_app(FOLDERED_SERVICE_NAME, config)
    check_plugin_installed(plugin_name, service_name=FOLDERED_SERVICE_NAME)

    config = sdk_marathon.get_config(FOLDERED_SERVICE_NAME)
    config['env']['TASKCFG_ALL_ELASTICSEARCH_PLUGINS'] = ""
    sdk_marathon.update_app(FOLDERED_SERVICE_NAME, config)
    check_plugin_uninstalled(plugin_name, service_name=FOLDERED_SERVICE_NAME)

@pytest.mark.recovery
@pytest.mark.sanity
def test_unchanged_scheduler_restarts_without_restarting_tasks():
    initial_task_ids = sdk_tasks.get_task_ids(FOLDERED_SERVICE_NAME, "master")
    shakedown.kill_process_on_host(sdk_marathon.get_scheduler_host(FOLDERED_SERVICE_NAME), "elastic.scheduler.Main")
    sdk_tasks.check_tasks_not_updated(FOLDERED_SERVICE_NAME, "master", initial_task_ids)

@pytest.mark.recovery
@pytest.mark.sanity
def test_bump_node_counts():
    # Run this test last, as it changes the task count
    config = sdk_marathon.get_config(FOLDERED_SERVICE_NAME)
    data_nodes = int(config['env']['DATA_NODE_COUNT'])
    config['env']['DATA_NODE_COUNT'] = str(data_nodes + 1)
    ingest_nodes = int(config['env']['INGEST_NODE_COUNT'])
    config['env']['INGEST_NODE_COUNT'] = str(ingest_nodes + 1)
    coordinator_nodes = int(config['env']['COORDINATOR_NODE_COUNT'])
    config['env']['COORDINATOR_NODE_COUNT'] = str(coordinator_nodes + 1)
    sdk_marathon.update_app(FOLDERED_SERVICE_NAME, config)
    sdk_tasks.check_running(FOLDERED_SERVICE_NAME, DEFAULT_TASK_COUNT + 3)
<commit_msg>Add missing import to test_shakedown<commit_after>import logging

import pytest

import sdk_cmd as cmd
import sdk_install
import sdk_metrics
import sdk_upgrade
import sdk_utils
from tests.config import *

log = logging.getLogger(__name__)

FOLDERED_SERVICE_NAME = sdk_utils.get_foldered_name(PACKAGE_NAME)

@pytest.fixture(scope='module', autouse=True)
def configure_package(configure_universe):
    try:
        log.info("Ensure elasticsearch and kibana are uninstalled...")
        sdk_install.uninstall(KIBANA_PACKAGE_NAME)
        sdk_install.uninstall(FOLDERED_SERVICE_NAME, package_name=PACKAGE_NAME)
        sdk_utils.gc_frameworks()

        sdk_upgrade.test_upgrade(
            "beta-{}".format(PACKAGE_NAME),
            PACKAGE_NAME,
            DEFAULT_TASK_COUNT,
            service_name=FOLDERED_SERVICE_NAME,
            additional_options={"service": {"name": FOLDERED_SERVICE_NAME}})

        yield  # let the test session execute
    finally:
        log.info("Clean up elasticsearch and kibana...")
        sdk_install.uninstall(KIBANA_PACKAGE_NAME)
        sdk_install.uninstall(FOLDERED_SERVICE_NAME, package_name=PACKAGE_NAME)

@pytest.fixture(autouse=True)
def pre_test_setup():
    sdk_tasks.check_running(FOLDERED_SERVICE_NAME, DEFAULT_TASK_COUNT)
    wait_for_expected_nodes_to_exist(service_name=FOLDERED_SERVICE_NAME)

@pytest.fixture
def default_populated_index():
    delete_index(DEFAULT_INDEX_NAME, service_name=FOLDERED_SERVICE_NAME)
    create_index(DEFAULT_INDEX_NAME, DEFAULT_SETTINGS_MAPPINGS, service_name=FOLDERED_SERVICE_NAME)
    create_document(DEFAULT_INDEX_NAME, DEFAULT_INDEX_TYPE, 1, {"name": "Loren", "role": "developer"}, service_name=FOLDERED_SERVICE_NAME)

@pytest.mark.sanity
@pytest.mark.smoke
def test_service_health():
    assert shakedown.service_healthy(FOLDERED_SERVICE_NAME)

@pytest.mark.sanity
def test_endpoints():
    # check that we can reach the scheduler via admin router, and that returned endpoints are sanitized:
    for endpoint in ENDPOINT_TYPES:
        endpoints = json.loads(cmd.run_cli('elastic --name={} endpoints {}'.format(FOLDERED_SERVICE_NAME, endpoint)))
        host = endpoint.split('-')[0] # 'coordinator-http' => 'coordinator'
        assert endpoints['dns'][0].startswith(sdk_hosts.autoip_host(FOLDERED_SERVICE_NAME, host + '-0-node'))
        assert endpoints['vip'].startswith(sdk_hosts.vip_host(FOLDERED_SERVICE_NAME, host))

@pytest.mark.sanity
def test_indexing(default_populated_index):
    indices_stats = get_elasticsearch_indices_stats(DEFAULT_INDEX_NAME, service_name=FOLDERED_SERVICE_NAME)
    assert indices_stats["_all"]["primaries"]["docs"]["count"] == 1
    doc = get_document(DEFAULT_INDEX_NAME, DEFAULT_INDEX_TYPE, 1, service_name=FOLDERED_SERVICE_NAME)
    assert doc["_source"]["name"] == "Loren"

@pytest.mark.sanity
@pytest.mark.metrics
@sdk_utils.dcos_1_9_or_higher
def test_metrics():
    sdk_metrics.wait_for_any_metrics(FOLDERED_SERVICE_NAME, "data-0-node", DEFAULT_ELASTIC_TIMEOUT)

@pytest.mark.focus
@pytest.mark.sanity
@pytest.mark.timeout(60 * 60)
def test_xpack_toggle_with_kibana(default_populated_index):
    log.info("\n***** Verify X-Pack disabled by default in elasticsearch")
    verify_commercial_api_status(False, service_name=FOLDERED_SERVICE_NAME)

    log.info("\n***** Test kibana with X-Pack disabled...")
    shakedown.install_package(KIBANA_PACKAGE_NAME, options_json={
        "kibana": {
            "elasticsearch_url": "http://" + sdk_hosts.vip_host(FOLDERED_SERVICE_NAME, "coordinator", 9200)
        }})
    shakedown.deployment_wait(app_id="/{}".format(KIBANA_PACKAGE_NAME), timeout=DEFAULT_KIBANA_TIMEOUT)
    check_kibana_adminrouter_integration("service/{}/".format(KIBANA_PACKAGE_NAME))
    log.info("Uninstall kibana with X-Pack disabled")
    sdk_install.uninstall(KIBANA_PACKAGE_NAME)

    log.info("\n***** Set/verify X-Pack enabled in elasticsearch")
    enable_xpack(service_name=FOLDERED_SERVICE_NAME)
    verify_commercial_api_status(True, service_name=FOLDERED_SERVICE_NAME)
    verify_xpack_license(service_name=FOLDERED_SERVICE_NAME)

    log.info("\n***** Write some data while enabled, disable X-Pack, and verify we can still read what we wrote.")
    create_document(
        DEFAULT_INDEX_NAME,
        DEFAULT_INDEX_TYPE,
        2,
        {"name": "X-Pack", "role": "commercial plugin"},
        service_name=FOLDERED_SERVICE_NAME)

    log.info("\n***** Test kibana with X-Pack enabled...")
    shakedown.install_package(KIBANA_PACKAGE_NAME, options_json={
        "kibana": {
            "elasticsearch_url": "http://" + sdk_hosts.vip_host(FOLDERED_SERVICE_NAME, "coordinator", 9200),
            "xpack_enabled": True
        }})
    log.info("\n***** Installing Kibana w/X-Pack can take as much as 15 minutes for Marathon deployment ")
    log.info("to complete due to a configured HTTP health check. (typical: 12 minutes)")
    shakedown.deployment_wait(app_id="/{}".format(KIBANA_PACKAGE_NAME), timeout=DEFAULT_KIBANA_TIMEOUT)
    check_kibana_adminrouter_integration("service/{}/login".format(KIBANA_PACKAGE_NAME))
    log.info("\n***** Uninstall kibana with X-Pack enabled")
    sdk_install.uninstall(KIBANA_PACKAGE_NAME)

    log.info("\n***** Disable X-Pack in elasticsearch.")
    disable_xpack(service_name=FOLDERED_SERVICE_NAME)
    log.info("\n***** Verify we can still read what we wrote when X-Pack was enabled.")
    verify_commercial_api_status(False, service_name=FOLDERED_SERVICE_NAME)
    doc = get_document(DEFAULT_INDEX_NAME, DEFAULT_INDEX_TYPE, 2, service_name=FOLDERED_SERVICE_NAME)
    assert doc["_source"]["name"] == "X-Pack"

@pytest.mark.recovery
@pytest.mark.sanity
def test_losing_and_regaining_index_health(default_populated_index):
    check_elasticsearch_index_health(DEFAULT_INDEX_NAME, "green", service_name=FOLDERED_SERVICE_NAME)
    shakedown.kill_process_on_host(sdk_hosts.system_host(FOLDERED_SERVICE_NAME, "data-0-node"), "data__.*Elasticsearch")
    check_elasticsearch_index_health(DEFAULT_INDEX_NAME, "yellow", service_name=FOLDERED_SERVICE_NAME)
    check_elasticsearch_index_health(DEFAULT_INDEX_NAME, "green", service_name=FOLDERED_SERVICE_NAME)

@pytest.mark.recovery
@pytest.mark.sanity
def test_master_reelection():
    initial_master = get_elasticsearch_master(service_name=FOLDERED_SERVICE_NAME)
    shakedown.kill_process_on_host(sdk_hosts.system_host(FOLDERED_SERVICE_NAME, initial_master), "master__.*Elasticsearch")
    wait_for_expected_nodes_to_exist(service_name=FOLDERED_SERVICE_NAME)
    new_master = get_elasticsearch_master(service_name=FOLDERED_SERVICE_NAME)
    assert new_master.startswith("master") and new_master != initial_master

@pytest.mark.recovery
@pytest.mark.sanity
def test_master_node_replace():
    # Ideally, the pod will get placed on a different agent. This test will verify that the remaining two masters
    # find the replaced master at its new IP address. This requires a reasonably low TTL for Java DNS lookups.
    cmd.run_cli('elastic --name={} pod replace master-0'.format(FOLDERED_SERVICE_NAME))
    # setup_function will verify that the cluster becomes healthy again.

@pytest.mark.recovery
@pytest.mark.sanity
def test_plugin_install_and_uninstall(default_populated_index):
    plugin_name = 'analysis-phonetic'
    config = sdk_marathon.get_config(FOLDERED_SERVICE_NAME)
    config['env']['TASKCFG_ALL_ELASTICSEARCH_PLUGINS'] = plugin_name
    sdk_marathon.update_app(FOLDERED_SERVICE_NAME, config)
    check_plugin_installed(plugin_name, service_name=FOLDERED_SERVICE_NAME)

    config = sdk_marathon.get_config(FOLDERED_SERVICE_NAME)
    config['env']['TASKCFG_ALL_ELASTICSEARCH_PLUGINS'] = ""
    sdk_marathon.update_app(FOLDERED_SERVICE_NAME, config)
    check_plugin_uninstalled(plugin_name, service_name=FOLDERED_SERVICE_NAME)

@pytest.mark.recovery
@pytest.mark.sanity
def test_unchanged_scheduler_restarts_without_restarting_tasks():
    initial_task_ids = sdk_tasks.get_task_ids(FOLDERED_SERVICE_NAME, "master")
    shakedown.kill_process_on_host(sdk_marathon.get_scheduler_host(FOLDERED_SERVICE_NAME), "elastic.scheduler.Main")
    sdk_tasks.check_tasks_not_updated(FOLDERED_SERVICE_NAME, "master", initial_task_ids)

@pytest.mark.recovery
@pytest.mark.sanity
def test_bump_node_counts():
    # Run this test last, as it changes the task count
    config = sdk_marathon.get_config(FOLDERED_SERVICE_NAME)
    data_nodes = int(config['env']['DATA_NODE_COUNT'])
    config['env']['DATA_NODE_COUNT'] = str(data_nodes + 1)
    ingest_nodes = int(config['env']['INGEST_NODE_COUNT'])
    config['env']['INGEST_NODE_COUNT'] = str(ingest_nodes + 1)
    coordinator_nodes = int(config['env']['COORDINATOR_NODE_COUNT'])
    config['env']['COORDINATOR_NODE_COUNT'] = str(coordinator_nodes + 1)
    sdk_marathon.update_app(FOLDERED_SERVICE_NAME, config)
    sdk_tasks.check_running(FOLDERED_SERVICE_NAME, DEFAULT_TASK_COUNT + 3)
<|endoftext|>


<commit_before>'''
Created on Oct 14, 2013

@author: Steven

This module contains numerous functions for fitting HOD models to correlation
function data. It uses MCMC techniques to do so.
'''
#===============================================================================
# IMPORTS
#===============================================================================
import numpy as np
import emcee
from hod import HOD
from scipy.stats import norm
import sys
from scipy.optimize import minimize
import copy
from multiprocessing import Pool, cpu_count
import scipy.sparse as sp
import scipy.sparse.linalg as spln
import time
#===============================================================================
# The Model
#===============================================================================

def model(parm, priors, h, attrs, data, sd, verbose):
    """
    Calculate the log probability of a HMF model given data

    Parameters
    ----------
    parm : sequence
        The position of the model. Takes arbitrary parameters.

    priors : list of prior classes
        A list containing objects that are either of the Uniform(), Normal() or
        MultiNorm() class. These specify the prior information on each parameter.

    h : ``hod.HOD`` instance
        A fully realised instance is best as it may make updating it faster

    attrs : list
        A list of the names of parameters passed in ``parm``. Corresponds to the
        names of the objects in ``priors``.

    data : array_like
        The measured correlation function.

    sd : array_like
        Uncertainty in the measured correlation function

    verbose : int
        How much to write to screen.

    Returns
    -------
    ll : float
        The log likelihood of the model at the given position.

    """
    ll = 0

    for prior in priors:
        # A uniform prior doesn't change likelihood but returns -inf if outside bounds
        if isinstance(prior, Uniform):
            index = attrs.index(prior.name)
            if parm[index] < prior.low or parm[index] > prior.high:
                ll = -np.inf

        elif isinstance(prior, Normal):
            index = attrs.index(prior.name)
            ll += norm.logpdf(parm[index], loc=prior.mean, scale=prior.sd)

        elif isinstance(prior, MultiNorm):
            indices = [attrs.index(name) for name in prior.name]
            ll += _lognormpdf(np.array(parm[indices]), np.array(prior.means),
                              prior.cov)

    # Rebuild the hod dict from given vals
    if not np.isinf(ll):
        hoddict = {attr:val for attr, val in zip(attrs, parm)}
        h.update(**hoddict)

        # The logprob of the model
        model = h.corr_gal.copy()

        ll += np.sum(norm.logpdf(data, loc=model, scale=sd))

    if verbose > 1:
        print parm
        print "Likelihood: ", ll
    return ll

def fit_hod(r, data, sd, priors, guess=[], nwalkers=100, nsamples=100, burnin=10,
           thin=50, nthreads=8, filename=None, chunks=None, verbose=0,
           **hodkwargs):
    """
    Run an MCMC procedure to fit a model correlation function to data

    Parameters
    ----------
    r : array_like
        The scales at which to perform analysis. Must be the same as the input
        data

    data : array_like
        The measured correlation function at ``r``

    sd : array_like
        The uncertainty in the measured correlation function

    priors : list of ``Uniform``, ``Gaussian`` or ``MultiGaussian`` objects
        A dictionary with keys as names of parameters in ``parm`` and values
        as a 4-item list. The first item is the guessed value for the parameter,
        the second is a string identifying whether the prior is normal (``norm``)
        or uniform (``unif``). The third and fourth are the uniform lower and
        upper boundaries or the mean and standard deviation of the normal
        distribution.

    guess : array_like, default ``None``
        Where to start the chain. If ``None``, will get central values from the
        distributions.

    nwalkers : int, default 100
        Number of walkers to use for Affine-Invariant Ensemble Sampler

    nsamples : int, default 100
        Number of samples that *each walker* will perform.

    burnin : int, default 10
        Number of samples from each walker that will be initially erased as burnin

    thin : int, default 50
        Keep 1 in every ``thin`` samples.

    nthreads : int, default 1
        Number of threads to use in sampling.

    filename : str, default ``None``
        A path to a file to which to write results sequentially

    chunks : int, default ``None``
        How many samples to run before appending results to file. Only
        applicable if ``filename`` is provided.

    verbose : int, default 0
        The verbosity level.

    \*\*hmfkwargs : arguments
        Any argument that could be sent to ``hmf.Perturbations``

    Returns
    -------
    flatchain : array_like, ``shape = (len(initial),nwalkers*nsamples)``
        The MCMC chain, with each parameter as a column

    acceptance_fraction : float
        The acceptance fraction for the MCMC run. Should be ....

    """

    if len(priors) == 0:
        raise ValueError("priors must be at least length 1")

    # Save which attributes are updatable for HOD as a list
    attrs = []
    for prior in priors:
        if isinstance(prior.name, basestring):
            attrs += [prior.name]
        else:
            attrs += prior.name

    # Get the number of variables for MCMC
    ndim = len(attrs)

    # Check that attrs are all applicable
    for a in attrs:
        if a not in ["wdm_mass", "delta_halo", "sigma_8", "n", "omegab", 'omegac',
                     "omegav", "omegak", "H0", "M_1", 'alpha', "M_min", 'gauss_width',
                     'M_0', 'fca', 'fcb', 'fs', 'delta', 'x', 'omegab_h2',
                     'omegac_h2', 'h']:
            raise ValueError(a + " is not a valid variable for MCMC in HOD")

    hodkwargs.update({"ThreadNum":nthreads})

    # Initialise the HOD object
    h = HOD(r=r, **hodkwargs)

    # It's better to get a corr_gal instance then the updates could be faster
    # BUT because of some reason we need to hack this and do it in a map() function
    h = Pool(1).apply(create_hod, [h])

    # re-set the number of threads used in pycamb to 1
    hodkwargs.update({"ThreadNum":1})

    # auto-calculate the number of threads to use if not set.
    if not nthreads:
        nthreads = cpu_count()

    # Set guess if not set
    if len(guess) != len(attrs):
        guess = []
        for prior in priors:
            if isinstance(prior, Uniform):
                guess += [(prior.high + prior.low) / 2]
            elif isinstance(prior, Normal):
                guess += [prior.mean]
            elif isinstance(prior, MultiNorm):
                guess += prior.means.tolist()

    guess = np.array(guess)

    # Get an initial value for all walkers, around a small ball near the initial guess
    stacked_val = guess.copy()
    for i in range(nwalkers - 1):
        stacked_val = np.vstack((guess, stacked_val))

    i = 0
    for prior in priors:
        if isinstance(prior, Uniform):
            stacked_val[:, i] += np.random.normal(loc=0.0, scale=0.05 * (prior.high - prior.low), size=nwalkers)
            i += 1
        elif isinstance(prior, Normal):
            stacked_val[:, i] += np.random.normal(loc=0.0, scale=prior.sd, size=nwalkers)
            i += 1
        elif isinstance(prior, MultiNorm):
            for j, name in enumerate(prior.name):
                stacked_val[:, i] += np.random.normal(loc=0.0, scale=np.sqrt(prior.cov[j, j]), size=nwalkers)
                i += 1

    sampler = emcee.EnsembleSampler(nwalkers, ndim, model,
                                    args=[priors, h, attrs, data, sd, verbose],
                                    threads=nthreads)

    if verbose:
        print attrs
    # Run a burn-in
    if burnin:
        pos, prob, state = sampler.run_mcmc(stacked_val, burnin)
        sampler.reset()
    else:
        pos = stacked_val

    # Run the actual run
    if filename is None:
        sampler.run_mcmc(pos, nsamples, thin=thin)
    else:
        header = "# " + "\t".join(attrs) + "\n"
        with open(filename, "w") as f:
            f.write(header)

        if chunks is None:
            chunks = nsamples

        start = time.time()
        for i, result in enumerate(sampler.sample(pos, iterations=nsamples)):
            if (i + 1) % chunks == 0:
                if verbose:
                    print "Done ", 100 * float(i + 1) / nsamples , "%. Time per sample: ", (time.time() - start) / ((i + 1) * nwalkers)
                with open(filename, "w") as f:
                    # need to write out nwalkers to be able to read the file back in properly
                    f.write("# %s\n" % (nwalkers))
                    f.write(header)
                    np.savetxt(f, sampler.flatchain[sampler.flatchain[:, 0] != 0.0, :])
    return sampler.flatchain, np.mean(sampler.acceptance_fraction)

class Uniform(object):
    def __init__(self, param, low, high):
        self.name = param
        self.low = low
        self.high = high

class Normal(object):
    def __init__(self, param, mean, sd):
        self.name = param
        self.mean = mean
        self.sd = sd

class MultiNorm(object):
    def __init__(self, params, means, cov):
        self.name = params
        self.means = means
        self.cov = cov

def create_hod(h):
    h.corr_gal
    return h

def _lognormpdf(x, mu, S):
    """ Calculate gaussian probability density of x, when x ~ N(mu,sigma) """
    nx = len(S)
    tmp = -0.5 * (nx * np.log(2 * np.pi) + np.linalg.slogdet(S)[1])

    err = x - mu
    if (sp.issparse(S)):
        numerator = spln.spsolve(S, err).T.dot(err)
    else:
        numerator = np.linalg.solve(S, err).T.dot(err)

    return tmp - numerator
<commit_msg>few more options to verbose<commit_after>'''
Created on Oct 14, 2013

@author: Steven

This module contains numerous functions for fitting HOD models to correlation
function data. It uses MCMC techniques to do so.
'''
#===============================================================================
# IMPORTS
#===============================================================================
import numpy as np
import emcee
from hod import HOD
from scipy.stats import norm
import sys
from scipy.optimize import minimize
import copy
from multiprocessing import Pool, cpu_count
import scipy.sparse as sp
import scipy.sparse.linalg as spln
import time
#===============================================================================
# The Model
#===============================================================================

def model(parm, priors, h, attrs, data, sd, verbose):
    """
    Calculate the log probability of a HMF model given data

    Parameters
    ----------
    parm : sequence
        The position of the model. Takes arbitrary parameters.

    priors : list of prior classes
        A list containing objects that are either of the Uniform(), Normal() or
        MultiNorm() class. These specify the prior information on each parameter.

    h : ``hod.HOD`` instance
        A fully realised instance is best as it may make updating it faster

    attrs : list
        A list of the names of parameters passed in ``parm``. Corresponds to the
        names of the objects in ``priors``.

    data : array_like
        The measured correlation function.

    sd : array_like
        Uncertainty in the measured correlation function

    verbose : int
        How much to write to screen.

    Returns
    -------
    ll : float
        The log likelihood of the model at the given position.

    """
    ll = 0

    for prior in priors:
        # A uniform prior doesn't change likelihood but returns -inf if outside bounds
        if isinstance(prior, Uniform):
            index = attrs.index(prior.name)
            if parm[index] < prior.low or parm[index] > prior.high:
                ll = -np.inf

        elif isinstance(prior, Normal):
            index = attrs.index(prior.name)
            ll += norm.logpdf(parm[index], loc=prior.mean, scale=prior.sd)

        elif isinstance(prior, MultiNorm):
            indices = [attrs.index(name) for name in prior.name]
            ll += _lognormpdf(np.array(parm[indices]), np.array(prior.means),
                              prior.cov)

    # Rebuild the hod dict from given vals
    if not np.isinf(ll):
        hoddict = {attr:val for attr, val in zip(attrs, parm)}
        h.update(**hoddict)

        # The logprob of the model
        model = h.corr_gal.copy()

        ll += np.sum(norm.logpdf(data, loc=model, scale=sd))

    if verbose > 1:
        print parm
        print data - model
        print norm.logpdf(data, loc=model, scale=sd)
        print "Likelihood: ", ll
    return ll

def fit_hod(r, data, sd, priors, guess=[], nwalkers=100, nsamples=100, burnin=10,
           thin=50, nthreads=8, filename=None, chunks=None, verbose=0,
           **hodkwargs):
    """
    Run an MCMC procedure to fit a model correlation function to data

    Parameters
    ----------
    r : array_like
        The scales at which to perform analysis. Must be the same as the input
        data

    data : array_like
        The measured correlation function at ``r``

    sd : array_like
        The uncertainty in the measured correlation function

    priors : list of ``Uniform``, ``Gaussian`` or ``MultiGaussian`` objects
        A dictionary with keys as names of parameters in ``parm`` and values
        as a 4-item list. The first item is the guessed value for the parameter,
        the second is a string identifying whether the prior is normal (``norm``)
        or uniform (``unif``). The third and fourth are the uniform lower and
        upper boundaries or the mean and standard deviation of the normal
        distribution.

    guess : array_like, default ``None``
        Where to start the chain. If ``None``, will get central values from the
        distributions.

    nwalkers : int, default 100
        Number of walkers to use for Affine-Invariant Ensemble Sampler

    nsamples : int, default 100
        Number of samples that *each walker* will perform.

    burnin : int, default 10
        Number of samples from each walker that will be initially erased as burnin

    thin : int, default 50
        Keep 1 in every ``thin`` samples.

    nthreads : int, default 1
        Number of threads to use in sampling.

    filename : str, default ``None``
        A path to a file to which to write results sequentially

    chunks : int, default ``None``
        How many samples to run before appending results to file. Only
        applicable if ``filename`` is provided.

    verbose : int, default 0
        The verbosity level.

    \*\*hmfkwargs : arguments
        Any argument that could be sent to ``hmf.Perturbations``

    Returns
    -------
    flatchain : array_like, ``shape = (len(initial),nwalkers*nsamples)``
        The MCMC chain, with each parameter as a column

    acceptance_fraction : float
        The acceptance fraction for the MCMC run. Should be ....

    """

    if len(priors) == 0:
        raise ValueError("priors must be at least length 1")

    # Save which attributes are updatable for HOD as a list
    attrs = []
    for prior in priors:
        if isinstance(prior.name, basestring):
            attrs += [prior.name]
        else:
            attrs += prior.name

    # Get the number of variables for MCMC
    ndim = len(attrs)

    # Check that attrs are all applicable
    for a in attrs:
        if a not in ["wdm_mass", "delta_halo", "sigma_8", "n", "omegab", 'omegac',
                     "omegav", "omegak", "H0", "M_1", 'alpha', "M_min", 'gauss_width',
                     'M_0', 'fca', 'fcb', 'fs', 'delta', 'x', 'omegab_h2',
                     'omegac_h2', 'h']:
            raise ValueError(a + " is not a valid variable for MCMC in HOD")

    hodkwargs.update({"ThreadNum":nthreads})

    # Initialise the HOD object
    h = HOD(r=r, **hodkwargs)

    # It's better to get a corr_gal instance then the updates could be faster
    # BUT because of some reason we need to hack this and do it in a map() function
    h = Pool(1).apply(create_hod, [h])

    # re-set the number of threads used in pycamb to 1
    hodkwargs.update({"ThreadNum":1})

    # auto-calculate the number of threads to use if not set.
    if not nthreads:
        nthreads = cpu_count()

    # Set guess if not set
    if len(guess) != len(attrs):
        guess = []
        for prior in priors:
            if isinstance(prior, Uniform):
                guess += [(prior.high + prior.low) / 2]
            elif isinstance(prior, Normal):
                guess += [prior.mean]
            elif isinstance(prior, MultiNorm):
                guess += prior.means.tolist()

    guess = np.array(guess)

    # Get an initial value for all walkers, around a small ball near the initial guess
    stacked_val = guess.copy()
    for i in range(nwalkers - 1):
        stacked_val = np.vstack((guess, stacked_val))

    i = 0
    for prior in priors:
        if isinstance(prior, Uniform):
            stacked_val[:, i] += np.random.normal(loc=0.0, scale=0.05 * (prior.high - prior.low), size=nwalkers)
            i += 1
        elif isinstance(prior, Normal):
            stacked_val[:, i] += np.random.normal(loc=0.0, scale=prior.sd, size=nwalkers)
            i += 1
        elif isinstance(prior, MultiNorm):
            for j, name in enumerate(prior.name):
                stacked_val[:, i] += np.random.normal(loc=0.0, scale=np.sqrt(prior.cov[j, j]), size=nwalkers)
                i += 1

    sampler = emcee.EnsembleSampler(nwalkers, ndim, model,
                                    args=[priors, h, attrs, data, sd, verbose],
                                    threads=nthreads)

    if verbose:
        print attrs
    # Run a burn-in
    if burnin:
        pos, prob, state = sampler.run_mcmc(stacked_val, burnin)
        sampler.reset()
    else:
        pos = stacked_val

    # Run the actual run
    if filename is None:
        sampler.run_mcmc(pos, nsamples, thin=thin)
    else:
        header = "# " + "\t".join(attrs) + "\n"
        with open(filename, "w") as f:
            f.write(header)

        if chunks is None:
            chunks = nsamples

        start = time.time()
        for i, result in enumerate(sampler.sample(pos, iterations=nsamples)):
            if (i + 1) % chunks == 0:
                if verbose:
                    print "Done ", 100 * float(i + 1) / nsamples , "%. Time per sample: ", (time.time() - start) / ((i + 1) * nwalkers)
                with open(filename, "w") as f:
                    # need to write out nwalkers to be able to read the file back in properly
                    f.write("# %s\n" % (nwalkers))
                    f.write(header)
                    np.savetxt(f, sampler.flatchain[sampler.flatchain[:, 0] != 0.0, :])
    return sampler.flatchain, np.mean(sampler.acceptance_fraction)

class Uniform(object):
    def __init__(self, param, low, high):
        self.name = param
        self.low = low
        self.high = high

class Normal(object):
    def __init__(self, param, mean, sd):
        self.name = param
        self.mean = mean
        self.sd = sd

class MultiNorm(object):
    def __init__(self, params, means, cov):
        self.name = params
        self.means = means
        self.cov = cov

def create_hod(h):
    h.corr_gal
    return h

def _lognormpdf(x, mu, S):
    """ Calculate gaussian probability density of x, when x ~ N(mu,sigma) """
    nx = len(S)
    tmp = -0.5 * (nx * np.log(2 * np.pi) + np.linalg.slogdet(S)[1])

    err = x - mu
    if (sp.issparse(S)):
        numerator = spln.spsolve(S, err).T.dot(err)
    else:
        numerator = np.linalg.solve(S, err).T.dot(err)

    return tmp - numerator
<|endoftext|>


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

import os
import time
import datetime
import ctypes
import re
import threading
import traceback
import atexit
from ctypes.wintypes import MAX_PATH

from logging import prnt

# [18:36:21.225] [Cartel Patrol Droid {2981965728841728}:3535188148330] [@Bellestarr] [Explosive Round {827176341471232}] [ApplyEffect {836045448945477}: Damage {836045448945501}] (1216 kinetic {836045448940873}) <1216>

# Global variables
parserThread = None

class GameEvent(object):
	TYPE_ABILITY_ACTIVATE   = 0x1
	TYPE_ABILITY_DEACTIVATE = 0x2
	TYPE_ABILITY_CANCEL     = 0x3
	TYPE_APPLY_BUFF         = 0x4
	TYPE_REMOVE_BUFF        = 0x5
	TYPE_ENTER_COMBAT       = 0x6
	TYPE_EXIT_COMBAT        = 0x7
	TYPE_DAMAGE             = 0x8
	TYPE_DEATH              = 0x9
	TYPE_HEAL               = 0x10

	def __init__(self):
		self.type = 0x0
		self.outgoing = False
		self.actor = None
		self.target = None
		self.ability = None
		self.abilityName = None
		self.actionType = None
		self.actionTypeName = None
		self.damage = 0
		self.healing = 0
		self.time = 0
		self.readTime = 0
		self.threat = 0
		self.enterEvent = False
		self.exitEvent = False
		self.inCombat = False
		self.recent = False

class Fight(object):
	def __init__(self):
		self.events = []
		self.enterEvent = None
		self.enterTime = 0
		self.exitEvent = None
		self.exitTime = 0

class Parser(object):
	"""docstring for Parser"""

	def __init__(self):
		self.linePat = re.compile("^\[(?P<hour>\d{1,2})\:(?P<minute>\d{2})\:(?P<second>\d{2})\.(?P<ms>\d{3})\] \[(?P<actor>[^\[\]]*)\] \[(?P<target>[^\[\]]*)\] \[(?:(?P<ability>[^{}]+))?(?: {(?P<abilityid>\d*)})?\] \[(?P<action>[^{}]+) {(?P<actionid>\d*)}: (?P<actiontype>[^{}]+) {(?P<actiontypeid>\d*)}\] \((?:(?P<result>[^\<\>]+))?\)(?: \<(?P<threat>-?\d*)\>)?$")
		self.logLocation = None
		self.fights = []
		self.ready = False
		self.me = None
		self.inCombat = False
		self.getDocs()
		pass

	def getDocs(self):
	 	dll = ctypes.windll.shell32
		buf = ctypes.create_unicode_buffer(MAX_PATH + 1)
		if dll.SHGetSpecialFolderPathW(None, buf, 0x0005, False):
		    self.logLocation = buf.value + "\\Star Wars - The Old Republic\\CombatLogs"
		    return True
		else:
		    return False

	def getNewestLog(self):
		if os.path.exists(self.logLocation):
			return max(os.listdir(self.logLocation))
		return None

	def run(self, hasStopped):
		prnt("Parser: Starting...")

		log = None
		try:
			if self.logLocation == None:
				raise Exception("No log location set. Did you forget to call getDocs?")

			logFile = self.getNewestLog()
			if logFile == None:
				prnt("Parser: Waiting for log...")
				while not hasStopped.isSet():
					logFile = self.getNewestLog()
					if logFile != None:
						break
					time.sleep(0.4)
				if hasStopped.isSet():
					return True
			log = open(self.logLocation + "\\" + logFile, 'r')
			logCursor = 0

			self.inCombat = False
			self.fights = []
			self.fight = None
			self.ready = False

			prnt("Parser: Began parsing %s"%logFile)

			lastLogFileCheck = time.time()

			while not hasStopped.isSet():
				if time.time() - lastLogFileCheck > 1:
					if logFile != self.getNewestLog():
						logFile = self.getNewestLog()
						prnt("Parser: Switched to parsing %s"%logFile)

						log.close()
						log = open(self.logLocation + "\\" + logFile, 'r')

						# Reset vars
						self.inCombat = False
						self.fights = []

						# Ensure we determine who "we" are again.
						self.me = None
					lastLogFileCheck = time.time()
				logCursor = log.tell()
				line = log.readline()
				if line == "":
					# Once we reach EOF mark us as ready for analyzation.
					if not self.ready:
						self.ready = True
					time.sleep(.25)
					continue

				res = self.linePat.match(line)
				if res:
					hour = int(res.group('hour'))
					minute = int(res.group('minute'))
					second = int(res.group('second'))
					ms = int(res.group('ms'))
					actor = res.group('actor')
					target = res.group('target')
					ability = res.group('ability')
					abilityId = res.group('abilityid')
					action = res.group('action')
					actionId = res.group('actionid')
					actionType = res.group('actiontype')
					actionTypeId = res.group('actiontypeid')
					result = res.group('result')
					threat = int(res.group('threat')) if res.group('threat') else 0

					today = time.mktime(datetime.datetime.now().replace(hour=0, minute=0, second=0, microsecond=0).timetuple())
					actionTime = today + (hour * 3600) + (minute * 60) + second + (ms / 1000.0)

					# Serious introspection here, man
					if self.me == None and actor == target and actor.find(':') == -1:
						prnt("Parser: Identified %s as me"%actor)
						self.me = actor

					event = GameEvent()
					event.type = self.resolveEventType(actionId, actionTypeId)
					event.actor = actor
					event.target = target
					event.ability = abilityId
					event.abilityName = ability
					event.actionType = actionTypeId
					event.actionTypeName = actionType
					event.inCombat = self.inCombat
					event.time = actionTime
					event.readTime = time.time()
					event.threat = threat
					event.enterEvent = False
					event.exitEvent = False
					event.recent = self.ready

					if event.type == GameEvent.TYPE_ENTER_COMBAT:
						event.enterEvent = True
					elif event.type == GameEvent.TYPE_EXIT_COMBAT:
						event.exitEvent = True
					elif event.type == GameEvent.TYPE_DEATH and self.inCombat and event.target == self.me:
						event.exitEvent = True
					elif event.type == GameEvent.TYPE_APPLY_BUFF and event.actionType == '973870949466372' and self.inCombat:
						# Safe Login Immunity
						event.exitEvent = True

					if event.enterEvent:
						fight = Fight()
						fight.enterEvent = event
						fight.enterTime = event.time
						self.fights.append(fight)
						self.fight = fight
						self.inCombat = True
					elif event.exitEvent:
						self.inCombat = False

					if event.type == GameEvent.TYPE_DAMAGE:
						sp = result.split(' ')
						if len(sp) > 2:
							dmg = sp[0]
							dmgType = sp[1]
							dmgTypeId = sp[2]
							if dmg.endswith('*'):
								dmg = dmg[:-1]
							dmg = int(dmg)
							if len(sp) == 6:
								fourth = sp[3]
								fifth = sp[4]
								sixth = sp[5]
								if fifth == 'absorbed' and event.target == self.me:
									absorbAmount = int(fourth[1:])
									dmg -= absorbAmount
						else:
							dmg = 0
						event.damage = dmg

					if event.type == GameEvent.TYPE_HEAL:
						heal = result
						if heal.endswith('*'):
							heal = heal[:-1]
						event.healing = int(heal)

					if self.fight != None:
						self.fight.events.append(event)

					if event.exitEvent and self.fight:
						self.fight.exitEvent = event
						self.fight.exitTime = event.time
						self.fight = None
				elif line[-1] != '\n' and line[-1] != '\r':
					prnt("Parser: Corrupted line! Backtracking to %d"%logCursor)
					log.seek(logCursor)
					time.sleep(0.2)

				time.sleep(.0001)
		except Exception:
			if log != None:
				log.close()
			prnt(traceback.format_exc())
			return False
		log.close()
		prnt("Parser: Exiting normally...")
		return True

	def resolveEventType(self, actionId, actionTypeId):
		if actionId == '836045448945472': # Event
			if actionTypeId == '836045448945479':
				return GameEvent.TYPE_ABILITY_ACTIVATE
			elif actionTypeId == '836045448945480':
				return GameEvent.TYPE_ABILITY_DEACTIVATE
			elif actionTypeId == '836045448945481':
				return GameEvent.TYPE_ABILITY_CANCEL
			elif actionTypeId == '836045448945489':
				return GameEvent.TYPE_ENTER_COMBAT
			elif actionTypeId == '836045448945490':
				return GameEvent.TYPE_EXIT_COMBAT
			elif actionTypeId == '836045448945493':
				return GameEvent.TYPE_DEATH
		elif actionId == '836045448945477': # ApplyEffect
			if actionTypeId == '836045448945501':
				return GameEvent.TYPE_DAMAGE
			elif actionTypeId == '836045448945500':
				return GameEvent.TYPE_HEAL
			else:
				return GameEvent.TYPE_APPLY_BUFF
		elif actionId == '836045448945478': # RemoveEffect
			return GameEvent.TYPE_REMOVE_BUFF
		return None

class ParserThread(threading.Thread):
	def __init__(self):
		threading.Thread.__init__(self)
		self.parser = None
		self.running = False
		self.setDaemon(True)

	def run(self):
		self.stopEvent = threading.Event()
		self.parser = Parser()
		crashCounter = 0
		while crashCounter < 6:
			self.running = True
			if not self.parser.run(self.stopEvent):
				prnt("ParserThread: Parser has crashed! Restarting... counter=%d"%crashCounter)
				crashCounter += 1
			else:
				break
			self.running = False

	def stop(self):
		self.stopEvent.set()

	def waitTillUp(self):
		while self.parser == None:
			time.sleep(0.0001)
			continue

	def isRunning(self):
		return self.running

	def getParser(self):
		return self.parser

def start():
	global parserThread
	parserThread = ParserThread()
	parserThread.start()

	atexit.register(shutdown)

def shutdown():
	global parserThread
	if parserThread != None:
		parserThread.stop()
		parserThread.join()
		parserThread = None

def get():
	return parserThread.parser

def getThread():
	return parserThread

if __name__ == '__main__':
	tParser = ParserThread()
	tParser.start()

	try:
		tParser.waitTillUp()
		parser = tParser.getParser()
		while tParser.isRunning():
			totalDamage = 0
			wasInCombat = False
			endTime = 0
			startTime = 0
			for ev in reversed(parser.events):
				if endTime == 0 and (ev.type == GameEvent.TYPE_EXIT_COMBAT or parser.events[-1].inCombat):
					wasInCombat = True
					endTime = ev.time
				if ev.type == GameEvent.TYPE_ENTER_COMBAT:
					startTime = ev.time
					break
				if ev.inCombat and ev.type == GameEvent.TYPE_DAMAGE:
					totalDamage += ev.damage
					wasInCombat = True

			tx = (endTime - startTime) / 1000.0
			if tx != 0:
				prnt("Damage dealt: %d, time: %d"%(totalDamage,tx))
				prnt("DPS: %d"%round(totalDamage / tx, 1))
			time.sleep(.01)
	except:
		prnt(traceback.format_exc())
		tParser.stop()
		tParser.join()
<commit_msg>New method for log times, and some cleanup<commit_after>#
# Copyright 2013 TeamSWAP
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import os
import time
import datetime
import ctypes
import re
import threading
import traceback
import atexit
from ctypes.wintypes import MAX_PATH

from logging import prnt

# [18:36:21.225] [Cartel Patrol Droid {2981965728841728}:3535188148330] [@Bellestarr] [Explosive Round {827176341471232}] [ApplyEffect {836045448945477}: Damage {836045448945501}] (1216 kinetic {836045448940873}) <1216>

# Global variables
parserThread = None

class GameEvent(object):
	TYPE_ABILITY_ACTIVATE   = 0x1
	TYPE_ABILITY_DEACTIVATE = 0x2
	TYPE_ABILITY_CANCEL     = 0x3
	TYPE_APPLY_BUFF         = 0x4
	TYPE_REMOVE_BUFF        = 0x5
	TYPE_ENTER_COMBAT       = 0x6
	TYPE_EXIT_COMBAT        = 0x7
	TYPE_DAMAGE             = 0x8
	TYPE_DEATH              = 0x9
	TYPE_HEAL               = 0x10

	def __init__(self):
		self.type = 0x0
		self.outgoing = False
		self.actor = None
		self.target = None
		self.ability = None
		self.abilityName = None
		self.actionType = None
		self.actionTypeName = None
		self.damage = 0
		self.healing = 0
		self.time = 0
		self.readTime = 0
		self.threat = 0
		self.enterEvent = False
		self.exitEvent = False
		self.inCombat = False
		self.recent = False

class Fight(object):
	def __init__(self):
		self.events = []
		self.enterEvent = None
		self.enterTime = 0
		self.exitEvent = None
		self.exitTime = 0

class Parser(object):
	"""docstring for Parser"""

	def __init__(self):
		self.linePat = re.compile("^\[(?P<hour>\d{1,2})\:(?P<minute>\d{2})\:(?P<second>\d{2})\.(?P<ms>\d{3})\] \[(?P<actor>[^\[\]]*)\] \[(?P<target>[^\[\]]*)\] \[(?:(?P<ability>[^{}]+))?(?: {(?P<abilityid>\d*)})?\] \[(?P<action>[^{}]+) {(?P<actionid>\d*)}: (?P<actiontype>[^{}]+) {(?P<actiontypeid>\d*)}\] \((?:(?P<result>[^\<\>]+))?\)(?: \<(?P<threat>-?\d*)\>)?$")
		self.logLocation = None
		self.fights = []
		self.ready = False
		self.me = None
		self.inCombat = False
		self.getDocs()
		pass

	def getDocs(self):
	 	dll = ctypes.windll.shell32
		buf = ctypes.create_unicode_buffer(MAX_PATH + 1)
		if dll.SHGetSpecialFolderPathW(None, buf, 0x0005, False):
		    self.logLocation = buf.value + "\\Star Wars - The Old Republic\\CombatLogs"
		    return True
		else:
		    return False

	def getNewestLog(self, onlyFilename=False):
		if os.path.exists(self.logLocation):
			filename = max(os.listdir(self.logLocation))
			fullPath = self.logLocation + "\\" + filename
			if onlyFilename:
				return filename
			return (filename, fullPath)
		return None

	def getMidnightTimestampForFile(self, filename):
		timestamp = os.path.getctime(filename)
		dt = datetime.datetime.fromtimestamp(timestamp)
		midnight = dt.replace(hour=0, minute=0, second=0, microsecond=0)
		midnightTimestamp = time.mktime(midnight.timetuple())
		return midnightTimestamp

	def run(self, hasStopped):
		prnt("Parser: Starting...")

		log = None
		try:
			if self.logLocation == None:
				raise Exception("No log location set. Did you forget to call getDocs?")

			logInfo = self.getNewestLog()
			if logInfo == None:
				prnt("Parser: Waiting for log...")
				while not hasStopped.isSet():
					logInfo = self.getNewestLog()
					if logInfo != None:
						break
					time.sleep(0.4)
				if hasStopped.isSet():
					return True
			(logFile, logPath) = logInfo

			log = open(logPath, 'r')

			self.inCombat = False
			self.fights = []
			self.fight = None
			self.ready = False

			logCursor = 0
			logDay = self.getMidnightTimestampForFile(logPath)
			logLastActionTime = 0

			prnt("Parser: Began parsing %s"%logFile)
			prnt("Parser: Log day is %s"%datetime.datetime.fromtimestamp(logDay))

			lastLogFileCheck = time.time()

			while not hasStopped.isSet():
				if time.time() - lastLogFileCheck > 1:
					if logFile != self.getNewestLog(onlyFilename=True):
						(logFile, logPath) = self.getNewestLog()

						prnt("Parser: Switched to parsing %s"%logFile)

						# Close previous log, and open new one.
						log.close()
						log = open(logPath, 'r')

						# Reset vars
						self.inCombat = False
						self.fights = []
						self.fight = None
						self.ready = False
						self.me = None

						logCursor = 0
						logDay = self.getMidnightTimestampForFile(logPath)
						logLastActionTime = 0
						prnt("Parser: Log day is %s"%datetime.datetime.fromtimestamp(logDay))
					lastLogFileCheck = time.time()
				logCursor = log.tell()
				line = log.readline()
				if line == "":
					# Once we reach EOF mark us as ready for analyzation.
					if not self.ready:
						self.ready = True
					time.sleep(.25)
					continue

				res = self.linePat.match(line)
				if res:
					hour = int(res.group('hour'))
					minute = int(res.group('minute'))
					second = int(res.group('second'))
					ms = int(res.group('ms'))
					actor = res.group('actor')
					target = res.group('target')
					ability = res.group('ability')
					abilityId = res.group('abilityid')
					action = res.group('action')
					actionId = res.group('actionid')
					actionType = res.group('actiontype')
					actionTypeId = res.group('actiontypeid')
					result = res.group('result')
					threat = int(res.group('threat')) if res.group('threat') else 0

					actionTime = logDay + (hour * 3600) + (minute * 60) + second + (ms / 1000.0)
					# Check for date rollover.
					if actionTime < logLastActionTime:
						logDay += 86400
						actionTime += 86400
						prnt("Parser: Rollover, day is now %s"%datetime.datetime.fromtimestamp(logDay))
					logLastActionTime = actionTime

					# Serious introspection here, man
					if self.me == None and actor == target and actor.find(':') == -1:
						prnt("Parser: Identified %s as me"%actor)
						self.me = actor

					event = GameEvent()
					event.type = self.resolveEventType(actionId, actionTypeId)
					event.actor = actor
					event.target = target
					event.ability = abilityId
					event.abilityName = ability
					event.actionType = actionTypeId
					event.actionTypeName = actionType
					event.inCombat = self.inCombat
					event.time = actionTime
					event.readTime = time.time()
					event.threat = threat
					event.enterEvent = False
					event.exitEvent = False
					event.recent = self.ready

					if event.type == GameEvent.TYPE_ENTER_COMBAT:
						event.enterEvent = True
					elif event.type == GameEvent.TYPE_EXIT_COMBAT:
						event.exitEvent = True
					elif event.type == GameEvent.TYPE_DEATH and self.inCombat and event.target == self.me:
						event.exitEvent = True
					elif event.type == GameEvent.TYPE_APPLY_BUFF and event.actionType == '973870949466372' and self.inCombat:
						# Safe Login Immunity
						event.exitEvent = True

					if event.enterEvent:
						fight = Fight()
						fight.enterEvent = event
						fight.enterTime = event.time
						self.fights.append(fight)
						self.fight = fight
						self.inCombat = True
					elif event.exitEvent:
						self.inCombat = False

					if event.type == GameEvent.TYPE_DAMAGE:
						sp = result.split(' ')
						if len(sp) > 2:
							dmg = sp[0]
							dmgType = sp[1]
							dmgTypeId = sp[2]
							if dmg.endswith('*'):
								dmg = dmg[:-1]
							dmg = int(dmg)
							if len(sp) == 6:
								fourth = sp[3]
								fifth = sp[4]
								sixth = sp[5]
								if fifth == 'absorbed' and event.target == self.me:
									absorbAmount = int(fourth[1:])
									dmg -= absorbAmount
						else:
							dmg = 0
						event.damage = dmg

					if event.type == GameEvent.TYPE_HEAL:
						heal = result
						if heal.endswith('*'):
							heal = heal[:-1]
						event.healing = int(heal)

					if self.fight != None:
						self.fight.events.append(event)

					if event.exitEvent and self.fight:
						self.fight.exitEvent = event
						self.fight.exitTime = event.time
						self.fight = None
				elif line[-1] != '\n' and line[-1] != '\r':
					prnt("Parser: Corrupted line! Backtracking to %d"%logCursor)
					log.seek(logCursor)
					time.sleep(0.2)

				time.sleep(.0001)
		except Exception:
			if log != None:
				log.close()
			prnt(traceback.format_exc())
			return False
		log.close()
		prnt("Parser: Exiting normally...")
		return True

	def resolveEventType(self, actionId, actionTypeId):
		if actionId == '836045448945472': # Event
			if actionTypeId == '836045448945479':
				return GameEvent.TYPE_ABILITY_ACTIVATE
			elif actionTypeId == '836045448945480':
				return GameEvent.TYPE_ABILITY_DEACTIVATE
			elif actionTypeId == '836045448945481':
				return GameEvent.TYPE_ABILITY_CANCEL
			elif actionTypeId == '836045448945489':
				return GameEvent.TYPE_ENTER_COMBAT
			elif actionTypeId == '836045448945490':
				return GameEvent.TYPE_EXIT_COMBAT
			elif actionTypeId == '836045448945493':
				return GameEvent.TYPE_DEATH
		elif actionId == '836045448945477': # ApplyEffect
			if actionTypeId == '836045448945501':
				return GameEvent.TYPE_DAMAGE
			elif actionTypeId == '836045448945500':
				return GameEvent.TYPE_HEAL
			else:
				return GameEvent.TYPE_APPLY_BUFF
		elif actionId == '836045448945478': # RemoveEffect
			return GameEvent.TYPE_REMOVE_BUFF
		return None

class ParserThread(threading.Thread):
	def __init__(self):
		threading.Thread.__init__(self)
		self.parser = None
		self.running = False
		self.setDaemon(True)

	def run(self):
		self.stopEvent = threading.Event()
		self.parser = Parser()
		crashCounter = 0
		while crashCounter < 6:
			self.running = True
			if not self.parser.run(self.stopEvent):
				prnt("ParserThread: Parser has crashed! Restarting... counter=%d"%crashCounter)
				crashCounter += 1
			else:
				break
			self.running = False

	def stop(self):
		self.stopEvent.set()

	def waitTillUp(self):
		while self.parser == None:
			time.sleep(0.0001)
			continue

	def isRunning(self):
		return self.running

	def getParser(self):
		return self.parser

def start():
	global parserThread
	parserThread = ParserThread()
	parserThread.start()

	atexit.register(shutdown)

def shutdown():
	global parserThread
	if parserThread != None:
		parserThread.stop()
		parserThread.join()
		parserThread = None

def get():
	return parserThread.parser

def getThread():
	return parserThread

if __name__ == '__main__':
	tParser = ParserThread()
	tParser.start()

	try:
		tParser.waitTillUp()
		parser = tParser.getParser()
		while tParser.isRunning():
			totalDamage = 0
			wasInCombat = False
			endTime = 0
			startTime = 0
			for ev in reversed(parser.events):
				if endTime == 0 and (ev.type == GameEvent.TYPE_EXIT_COMBAT or parser.events[-1].inCombat):
					wasInCombat = True
					endTime = ev.time
				if ev.type == GameEvent.TYPE_ENTER_COMBAT:
					startTime = ev.time
					break
				if ev.inCombat and ev.type == GameEvent.TYPE_DAMAGE:
					totalDamage += ev.damage
					wasInCombat = True

			tx = (endTime - startTime) / 1000.0
			if tx != 0:
				prnt("Damage dealt: %d, time: %d"%(totalDamage,tx))
				prnt("DPS: %d"%round(totalDamage / tx, 1))
			time.sleep(.01)
	except:
		prnt(traceback.format_exc())
		tParser.stop()
		tParser.join()
<|endoftext|>


<commit_before>#!/usr/bin/env python
import socket
import os
import subprocess
import csv
import sys
import re
import time
import argparse
import shlex

# Override builtin: raw_input was renamed to input in python3 (PEP 3111)
try:
    input = raw_input
except NameError:
    pass

def create_output_folder_name():
    """
    Return the output folder name in format YYYY-MM-DD_output.

    :return:
    """
    date_today = time.strftime("%Y-%m-%d")
    folder_name = "%s_output" % date_today
    return folder_name

def create_csv_file_name():
    """
    Return the CSV file name in format "Hr_Min_{AM/PM}.csv".

    :return:
    """
    time_now = time.strftime("%I_%M_%p")
    csv_file_name = "%s.csv" % time_now
    if csv_file_name[0] == "0":  # strip leading zero off
        csv_file_name = csv_file_name.lstrip('0')
    return csv_file_name

# Create the output folder name
folder_name = create_output_folder_name()

# Create the CSV file name
csv_file_name = create_csv_file_name()

def detect_os():
    """
    Detect OS of system and sets os commands for use in various functions.

    :return: dict()
    """
    os_info = {}
    if os.name == 'nt':  # Windows
        os_info['os'] = "win"
        os_info['arp_cmd'] = "arp -a {ip}"
        os_info['ping_cmd'] = "ping -n 1 {ip}"
    if os.name == 'posix':  # posix
        os_info['os'] = "posix"
        os_info['arp_cmd'] = "arp -a | grep -w {ip}"
        os_info['ping_cmd'] = "ping -c 1 {ip}"
    # check for OSX last because it also shows up as posix in os.name
    if sys.platform == 'darwin':  # OSX
        os_info['os'] = "osx"
        os_info['arp_cmd'] = "arp -a | grep -w {ip}"
        os_info['ping_cmd'] = "ping -c 1 {ip}"
        os_info['smbutil'] = "smbutil status {ip} | grep Server"
    return os_info

# Set OS information
detected_os = detect_os()

def subproc_pipe_runner(ip, command):
    """
    Take an ip and a "raw" shell input command. Splits it on "|" if present,
    and runs it without dependency on shell=True.

    :param ip: The IP to run the command against
    :param command: The shell command to interpret. Max one pipe ("|")
    :return: string
    """
    command = command.replace("{ip}", "{0}")
    command = command.format(ip)

    #  Get both piped and unpiped commands into the same format:
    #  a list with one value
    if "|" in command:
        new_command = command.split("|")
    else:
        new_command = []
        new_command.append(command)

    if len(new_command) == 1:
        process_one = subprocess.\
            Popen(shlex.split("{0}".format(new_command[0])),
                  stdout=subprocess.PIPE)
        output = process_one.communicate()
        return output
    if len(new_command) == 2:
        process_one = subprocess.\
            Popen(shlex.split("{0}".format(new_command[0])),
                  stdout=subprocess.PIPE)
        process_two = subprocess.\
            Popen(shlex.split("{0}".format(new_command[1])),
                  stdin=process_one.stdout, stdout=subprocess.PIPE)
        process_one.stdout.close()
        output = process_two.communicate()
        return output

def make_dir(folder_name):
    """
    Create the output folder name if it doesn't exist yet.

    :param folder_name:
    :return:
    """
    if not os.path.exists("./%s" % folder_name):
        try:
            os.makedirs("./%s" % folder_name)
            print("Created output folder /%s" % folder_name)
        except OSError as err:
            print("Unable to create folder /%s" % folder_name)
            exists_error = re.search("exists", str(err))
            if exists_error:  # in theory we should never trigger this
                print("Reason: Folder already exists.")
            perms_error = re.search("denied", str(err))
            if perms_error:
                print("Reason: Insufficient permissions.")
            if not exists_error and not perms_error:
                print("Encountered error while creating output folder:")
                print(err)

def ip_check(input_ip):
    """
    Verify correct format of input_ip: "xxx.xxx.xxx.xxx".

    :param input_ip: string
    :return:
    """
    ipregex = '\\b(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.)' \
              '{3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\\b'
    validIP = re.match(ipregex, input_ip)
    if validIP:
        return True
    else:
        return False

def nslooky(ip, detected_os):
    """
    Given an IP, returns hostname.

    :param ip: string
    :param detected_os: dict()
    :return: string
    """
    try:
        output = socket.gethostbyaddr(ip)
        return output[0]
    except Exception as err:
        if detected_os['os'] == 'osx':
            output = subproc_pipe_runner(ip, detected_os["smbutil"])
            if output[0] == "":
                output = "No hostname found"
                return output
            output = output[0].split(' ')[1].strip()
            return output
        not_found_error = re.search("not found", str(err))
        if not_found_error:  # Windows: catch [Errno 11004] host not found
            output = "No hostname found"
        else:
            output = "No hostname found"
        return output

def get_name(ip):
    """
    Return the name of pinged host.

    :param ip: string
    :return: string
    """
    try:
        name = nslooky(ip, detected_os)
    # We shouldn't ever hit this, should catch this in nslooky()
    except Exception as err:
        name = "Gen. except error in get_name()"
    return name

def title_check(folder_name, csv_file_name):
    """
    Create titles (headers) in .csv output file.

    :param folder_name: string
    :param csv_file_name: string
    :return:
    """
    exists = os.path.exists(r"./{0}/{1}".format(folder_name, csv_file_name))
    if exists is True:
        pass
    if exists is False:
        make_dir(folder_name)
        myfile = open("./{0}/{1}".format(folder_name, csv_file_name), "w")
        wr = csv.writer(myfile)
        titles = ["ip", "ping ms time", "hostname", "mac"]
        wr.writerow(titles)

def get_ping_ms_response(ip, detected_os):
    """
    Return ping ms response time of pinged host.

    :param ip: string
    :param detected_os: dict()
    :return:
    """
    pingResult = subproc_pipe_runner(ip, detected_os['ping_cmd'])
    ping_found = re.search(r'time[=<]?(\d*[\.]?\d*\s?ms)?', str(pingResult[0]))
    if ping_found and ping_found.group(1):
        ping_msResponse = ping_found.group(1)
    else:
        ping_msResponse = 'Host unreachable'
    return ping_msResponse

def get_mac(ip, detected_os):
    """
    Given an IP, return a MAC address.

    :param ip: string
    :param detected_os: string
    :return: string
    """
    arpResult = subproc_pipe_runner(ip, detected_os['arp_cmd'])
    find_mac = re.search(r'[\b\s]*(([0-9A-F]{2}[:-]){5}([0-9A-F]{2}))[\b\s]*',
                         str(arpResult[0].upper()))
    if find_mac:
        item = re.sub('-', ':', find_mac.group(1))
    else:
        item = 'No MAC addr. found'
    return item

def get_results(ip, folder_name, csv_file_name,
                start=1, end=255, get_all=False, csv_out=False):
    """
    Given inputs, get and display query results to end user.
    Writes CSV if applicable.

    :param ip: string
    :param folder_name: string
    :param csv_file_name: string
    :param start: int
    :param end: int
    :param get_all: boolean
    :param csv_out: boolean
    :return:
    """
    if csv_out:
        try:
            myfile = open("./%s/%s" % (folder_name, csv_file_name), "a")
            wr = csv.writer(myfile)
        except IOError as e:
            print("Could not open /{2}/{3}: I/O error({0}): {1}".format(
                  e.errno, e.strerror, folder_name, csv_file_name))
            sys.exit()
    print("\nResults:")
    if get_all:
        first_three = re.match(r'((\d{,3}\.\d{,3}\.\d{,3})\.)?(\d{,3})', ip)
        if int(end) <= 254:
            end += 1
        for address in range(int(start), int(end)):
            ip = first_three.group(1) + str(address)
            ping = getPing_msResponse(ip, detected_os)
            name = nslooky(ip, detected_os)
            mac = getMac(ip, detected_os)
            if csv_out:
                wr.writerow([ip, ping, name, mac])
            print("{0} {1} {2} {3}".format(ip, ping, name, mac))
    else:
        ping = getPing_msResponse(ip, detected_os)
        name = nslooky(ip, detected_os)
        mac = getMac(ip, detected_os)
        if csv_out:
            wr.writerow([ip, ping, name, mac])
        print("{0} {1} {2} {3}".format(ip, ping, name, mac))
    if csv_out:
        myfile.close()

def detect_ip(ip_address=None):
    """
    Create a UDP socket connection to populate getsockname().
    The address does not actually need to resolve ie:
    Google DNS 8.8.8.8 used here.

    :param ip_address: string
    :return: string
    """
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(('8.8.8.8', 80))
        ip_address = s.getsockname()[0]
    except socket.error:
        print("Failed to detect IP of current host!")
        choice = input("(I)nput host IP manually, or (Q)uit?: ")
        if choice.upper() == "I":
            while True:
                if not ip_address or not ip_check(ip_address):
                    ip_address = input("INPUT IP: ")
                else:
                    break
        else:
            sys.exit("\nQuitting..")
    finally:
        s.close()
    return ip_address

def main(ip=None, start=1, end=255, get_all=False, csv_out=False):
    """
    Main program body.

    :param ip: string
    :param start: int
    :param end: int
    :param get_all: boolean
    :param csv_out: boolean
    :return:
    """
    answer = int()
    ip_range = sorted([int(start), int(end)])  # ensure range from low to high
    while answer != 3:
        print('\n\nQuery based on IP: {0}\n'
              '1) Continue with detected IP (entries for x.x.x.{1}-{2})\n'
              '2) Enter another IP (creates one entry)\n'
              '3) Exit'.format(ip, ip_range[0], int(ip_range[1])))
        try:
            if get_all:
                answer = 1
            else:
                answer = int(input("Selection? "))
            if answer == 1:
                get_all = True
                print('\n')
                if csv_out:
                    title_check(folder_name, csv_file_name)
                get_results(ip, folder_name, csv_file_name,
                            end=ip_range[1], start=ip_range[0],
                            get_all=get_all, csv_out=csv_out)
                get_all = False  # Clear the get_all flag to break loop
            elif answer == 2:
                choiceIP = input("Input IP: ")
                if ip_check(choiceIP):
                    if csv_out:
                        title_check(folder_name, csv_file_name)
                    get_results(choiceIP, folder_name, csv_file_name,
                                start=ip_range[0], end=ip_range[1],
                                get_all=get_all, csv_out=csv_out)
                else:
                    print("Invalid IP")
            elif answer == 3:
                sys.exit()
        except ValueError:
            print("\n")
            print("Invalid entry")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="HostMAC")
    parser.add_argument('-csv', help='log output to csv', action='store_true',
                        required=False)
    parser.add_argument('-ip', default=detect_ip(),
                        help='specify ip, default: current ip', required=False)
    parser.add_argument('-all', help='check range 1-254', action='store_true',
                        required=False)
    parser.add_argument('-start', default=1, type=int, choices=range(1, 254),
                        metavar='', help='start of range', required=False)
    parser.add_argument('-end', default=255, type=int, choices=range(2, 255),
                        metavar='', help='end of range', required=False)
    args = parser.parse_args()
    try:
        main(ip=args.ip, start=args.start, end=args.end,
             get_all=args.all, csv_out=args.csv)
    except KeyboardInterrupt:
        print("\nCanceled by user.. Exiting.")
<commit_msg>Closes #59 and #60<commit_after>#!/usr/bin/env python
import socket
import os
import subprocess
import csv
import sys
import re
import time
import argparse
import shlex

# Override builtin: raw_input was renamed to input in python3 (PEP 3111)
try:
    input = raw_input
except NameError:
    pass

def create_output_folder_name():
    """
    Return the output folder name in format YYYY-MM-DD_output.

    :return:
    """
    date_today = time.strftime("%Y-%m-%d")
    folder_name = "%s_output" % date_today
    return folder_name

def create_csv_file_name():
    """
    Return the CSV file name in format "Hr_Min_{AM/PM}.csv".

    :return:
    """
    time_now = time.strftime("%I_%M_%p")
    csv_file_name = "%s.csv" % time_now
    if csv_file_name[0] == "0":  # strip leading zero off
        csv_file_name = csv_file_name.lstrip('0')
    return csv_file_name

# Create the output folder name
folder_name = create_output_folder_name()

# Create the CSV file name
csv_file_name = create_csv_file_name()

def detect_os():
    """
    Detect OS of system and sets os commands for use in various functions.

    :return: dict()
    """
    os_info = {}
    if os.name == 'nt':  # Windows
        os_info['os'] = "win"
        os_info['arp_cmd'] = "arp -a {ip}"
        os_info['ping_cmd'] = "ping -n 1 {ip}"
    if os.name == 'posix':  # posix
        os_info['os'] = "posix"
        os_info['arp_cmd'] = "arp -a | grep -w {ip}"
        os_info['ping_cmd'] = "ping -c 1 {ip}"
    # check for OSX last because it also shows up as posix in os.name
    if sys.platform == 'darwin':  # OSX
        os_info['os'] = "osx"
        os_info['arp_cmd'] = "arp -a | grep -w {ip}"
        os_info['ping_cmd'] = "ping -c 1 {ip}"
        os_info['smbutil'] = "smbutil status {ip} | grep Server"
    return os_info

# Set OS information
detected_os = detect_os()

def subproc_pipe_runner(ip, command):
    """
    Take an ip and a "raw" shell input command. Splits it on "|" if present,
    and runs it without dependency on shell=True.

    :param ip: The IP to run the command against
    :param command: The shell command to interpret. Max one pipe ("|")
    :return: string
    """
    command = command.replace("{ip}", "{0}")
    command = command.format(ip)

    #  Get both piped and unpiped commands into the same format:
    #  a list with one value
    if "|" in command:
        new_command = command.split("|")
    else:
        new_command = []
        new_command.append(command)

    if len(new_command) == 1:
        process_one = subprocess.\
            Popen(shlex.split("{0}".format(new_command[0])),
                  stdout=subprocess.PIPE)
        output = process_one.communicate()
        return output
    if len(new_command) == 2:
        process_one = subprocess.\
            Popen(shlex.split("{0}".format(new_command[0])),
                  stdout=subprocess.PIPE)
        process_two = subprocess.\
            Popen(shlex.split("{0}".format(new_command[1])),
                  stdin=process_one.stdout, stdout=subprocess.PIPE)
        process_one.stdout.close()
        output = process_two.communicate()
        return output

def make_dir(folder_name):
    """
    Create the output folder name if it doesn't exist yet.

    :param folder_name:
    :return:
    """
    if not os.path.exists("./%s" % folder_name):
        try:
            os.makedirs("./%s" % folder_name)
            print("Created output folder /%s" % folder_name)
        except OSError as err:
            print("Unable to create folder /%s" % folder_name)
            exists_error = re.search("exists", str(err))
            if exists_error:  # in theory we should never trigger this
                print("Reason: Folder already exists.")
            perms_error = re.search("denied", str(err))
            if perms_error:
                print("Reason: Insufficient permissions.")
            if not exists_error and not perms_error:
                print("Encountered error while creating output folder:")
                print(err)

def ip_check(input_ip):
    """
    Verify correct format of input_ip: "xxx.xxx.xxx.xxx".

    :param input_ip: string
    :return:
    """
    ipregex = '\\b(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.)' \
              '{3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\\b'
    validIP = re.match(ipregex, input_ip)
    if validIP:
        return True
    else:
        return False

def nslooky(ip, detected_os):
    """
    Given an IP, returns hostname.

    :param ip: string
    :param detected_os: dict()
    :return: string
    """
    try:
        output = socket.gethostbyaddr(ip)
        return output[0]
    except Exception as err:
        if detected_os['os'] == 'osx':
            output = subproc_pipe_runner(ip, detected_os["smbutil"])
            if output[0] == "":
                output = "No hostname found"
                return output
            output = output[0].split(' ')[1].strip()
            return output
        not_found_error = re.search("not found", str(err))
        if not_found_error:  # Windows: catch [Errno 11004] host not found
            output = "No hostname found"
        else:
            output = "No hostname found"
        return output

def get_name(ip):
    """
    Return the name of pinged host.

    :param ip: string
    :return: string
    """
    try:
        name = nslooky(ip, detected_os)
    # We shouldn't ever hit this, should catch this in nslooky()
    except Exception as err:
        name = "Gen. except error in get_name()"
    return name

def title_check(folder_name, csv_file_name):
    """
    Create titles (headers) in .csv output file.

    :param folder_name: string
    :param csv_file_name: string
    :return:
    """
    exists = os.path.exists(r"./{0}/{1}".format(folder_name, csv_file_name))
    if exists is True:
        pass
    if exists is False:
        make_dir(folder_name)
        myfile = open("./{0}/{1}".format(folder_name, csv_file_name), "w")
        wr = csv.writer(myfile)
        titles = ["ip", "ping ms time", "hostname", "mac"]
        wr.writerow(titles)

def get_ping_ms_response(ip, detected_os):
    """
    Return ping ms response time of pinged host.

    :param ip: string
    :param detected_os: dict()
    :return:
    """
    pingResult = subproc_pipe_runner(ip, detected_os['ping_cmd'])
    ping_found = re.search(r'time[=<]?(\d*[\.]?\d*\s?ms)?', str(pingResult[0]))
    if ping_found and ping_found.group(1):
        ping_msResponse = ping_found.group(1)
    else:
        ping_msResponse = 'Host unreachable'
    return ping_msResponse

def get_mac(ip, detected_os):
    """
    Given an IP, return a MAC address.

    :param ip: string
    :param detected_os: string
    :return: string
    """
    arpResult = subproc_pipe_runner(ip, detected_os['arp_cmd'])
    find_mac = re.search(r'[\b\s]*(([0-9A-F]{2}[:-]){5}([0-9A-F]{2}))[\b\s]*',
                         str(arpResult[0].upper()))
    if find_mac:
        item = re.sub('-', ':', find_mac.group(1))
    else:
        item = 'No MAC addr. found'
    return item

def get_results(ip, folder_name, csv_file_name,
                start=1, end=255, get_all=False, csv_out=False):
    """
    Given inputs, get and display query results to end user.
    Writes CSV if applicable.

    :param ip: string
    :param folder_name: string
    :param csv_file_name: string
    :param start: int
    :param end: int
    :param get_all: boolean
    :param csv_out: boolean
    :return:
    """
    if csv_out:
        try:
            myfile = open("./%s/%s" % (folder_name, csv_file_name), "a")
            wr = csv.writer(myfile)
        except IOError as e:
            print("Could not open /{2}/{3}: I/O error({0}): {1}".format(
                  e.errno, e.strerror, folder_name, csv_file_name))
            sys.exit()
    print("\nResults:")
    if get_all:
        first_three = re.match(r'((\d{,3}\.\d{,3}\.\d{,3})\.)?(\d{,3})', ip)
        if int(end) <= 254:
            end += 1
        for address in range(int(start), int(end)):
            ip = first_three.group(1) + str(address)
            ping = get_ping_ms_response(ip, detected_os)
            name = nslooky(ip, detected_os)
            mac = get_mac(ip, detected_os)
            if csv_out:
                wr.writerow([ip, ping, name, mac])
            print("{0} {1} {2} {3}".format(ip, ping, name, mac))
    else:
        ping = get_ping_ms_response(ip, detected_os)
        name = nslooky(ip, detected_os)
        mac = get_mac(ip, detected_os)
        if csv_out:
            wr.writerow([ip, ping, name, mac])
        print("{0} {1} {2} {3}".format(ip, ping, name, mac))
    if csv_out:
        myfile.close()

def detect_ip(ip_address=None):
    """
    Create a UDP socket connection to populate getsockname().
    The address does not actually need to resolve ie:
    Google DNS 8.8.8.8 used here.

    :param ip_address: string
    :return: string
    """
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(('8.8.8.8', 80))
        ip_address = s.getsockname()[0]
    except socket.error:
        print("Failed to detect IP of current host!")
        choice = input("(I)nput host IP manually, or (Q)uit?: ")
        if choice.upper() == "I":
            while True:
                if not ip_address or not ip_check(ip_address):
                    ip_address = input("INPUT IP: ")
                else:
                    break
        else:
            sys.exit("\nQuitting..")
    finally:
        s.close()
    return ip_address

def main(ip=None, start=1, end=255, get_all=False, csv_out=False):
    """
    Main program body.

    :param ip: string
    :param start: int
    :param end: int
    :param get_all: boolean
    :param csv_out: boolean
    :return:
    """
    answer = int()
    ip_range = sorted([int(start), int(end)])  # ensure range from low to high
    while answer != 3:
        print('\n\nQuery based on IP: {0}\n'
              '1) Continue with detected IP (entries for x.x.x.{1}-{2})\n'
              '2) Enter another IP (creates one entry)\n'
              '3) Exit'.format(ip, ip_range[0], int(ip_range[1])))
        try:
            if get_all:
                answer = 1
            else:
                answer = int(input("Selection? "))
            if answer == 1:
                get_all = True
                print('\n')
                if csv_out:
                    title_check(folder_name, csv_file_name)
                get_results(ip, folder_name, csv_file_name,
                            end=ip_range[1], start=ip_range[0],
                            get_all=get_all, csv_out=csv_out)
                get_all = False  # Clear the get_all flag to break loop
            elif answer == 2:
                choiceIP = input("Input IP: ")
                if ip_check(choiceIP):
                    if csv_out:
                        title_check(folder_name, csv_file_name)
                    get_results(choiceIP, folder_name, csv_file_name,
                                start=ip_range[0], end=ip_range[1],
                                get_all=get_all, csv_out=csv_out)
                else:
                    print("Invalid IP")
            elif answer == 3:
                sys.exit()
        except ValueError:
            print("\n")
            print("Invalid entry")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="HostMAC")
    parser.add_argument('-csv', help='log output to csv', action='store_true',
                        required=False)
    parser.add_argument('-ip', default=detect_ip(),
                        help='specify ip, default: current ip', required=False)
    parser.add_argument('-all', help='check range 1-254', action='store_true',
                        required=False)
    parser.add_argument('-start', '-s', default=1, type=int, choices=range(1, 254),
                        metavar='', help='start of range', required=False)
    parser.add_argument('-end', '-e', default=255, type=int, choices=range(2, 255),
                        metavar='', help='end of range', required=False)
    args = parser.parse_args()
    try:
        main(ip=args.ip, start=args.start, end=args.end,
             get_all=args.all, csv_out=args.csv)
    except KeyboardInterrupt:
        print("\nCanceled by user.. Exiting.")
<|endoftext|>


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

# Pack a webpage including images and CSS into a single HTML file.

import sys
import requests
import base64
from urllib.parse import urljoin
from bs4 import BeautifulSoup

#TODO: Properly parse command line

pageurl = sys.argv[1]
print("Processing {}".format(pageurl))

soup = BeautifulSoup(requests.get(pageurl).text)
imgtags = soup.find_all('img')
for image in imgtags:
    print("Image found: " + image['src'])
    url = urljoin(pageurl, image['src'])
    imagecontent = requests.get(url).content
    imageencoded = base64.b64encode(imagecontent).decode()
    image['src'] = "data:image/jpeg;base64," + imageencoded

# Conversion complete, write output and cleanup
outfile = open('out.html', 'w')
outfile.write(str(soup))
outfile.close()
print("All done, file written to " + "out.html")
<commit_msg>Moved code into main() function<commit_after>#!/usr/bin/env python3

# Pack a webpage including images and CSS into a single HTML file.

import sys
import requests
import base64
from urllib.parse import urljoin
from bs4 import BeautifulSoup

#TODO: Properly parse command line
#TODO: Ignore files already data-URI encoded

def main():
    pageurl = sys.argv[1]
    print("Processing {}".format(pageurl))

    soup = BeautifulSoup(requests.get(pageurl).text)
    imgtags = soup.find_all('img')
    for image in imgtags:
        print("Image found: " + image['src'])
        url = urljoin(pageurl, image['src'])
        imagecontent = requests.get(url).content
        imageencoded = base64.b64encode(imagecontent).decode()
        image['src'] = "data:image/jpeg;base64," + imageencoded

    # Conversion complete, write output and cleanup
    outfile = open('out.html', 'w')
    outfile.write(str(soup))
    outfile.close()
    print("All done, file written to " + "out.html")

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