<commit_before>from __future__ import (absolute_import, unicode_literals, division,
                        print_function)
import numpy as np
import logging
import warnings
import os

from astropy.io import fits
from astropy.table import Table

import stingray.utils as utils
from .utils import order_list_of_arrays, is_string
from .utils import assign_value_if_none

from .gti import  _get_gti_from_extension, load_gtis

try:
    # Python 2
    import cPickle as pickle
except:
    # Python 3
    import pickle

def get_file_extension(fname):
    """Get the extension from the file name."""
    return os.path.splitext(fname)[1]

def high_precision_keyword_read(hdr, keyword):
    """Read FITS header keywords, also if split in two.

    In the case where the keyword is split in two, like

        MJDREF = MJDREFI + MJDREFF

    in some missions, this function returns the summed value. Otherwise, the
    content of the single keyword

    Parameters
    ----------
    hdr : dict_like
        The FITS header structure, or a dictionary
    keyword : str
        The key to read in the header

    Returns
    -------
    value : long double
        The value of the key, or None if something went wrong

    """
    try:
        value = np.longdouble(hdr[keyword])
        return value
    except:
        pass
    try:
        if len(keyword) == 8:
            keyword = keyword[:7]
        value = np.longdouble(hdr[keyword + 'I'])
        value += np.longdouble(hdr[keyword + 'F'])
        return value
    except:
        return None

def _get_additional_data(lctable, additional_columns):
    additional_data = {}
    if additional_columns is not None:
        for a in additional_columns:
            try:
                additional_data[a] = np.array(lctable.field(a))
            except:  # pragma: no cover
                if a == 'PI':
                    logging.warning('Column PI not found. Trying with PHA')
                    additional_data[a] = np.array(lctable.field('PHA'))
                else:
                    raise Exception('Column' + a + 'not found')

    return additional_data

def load_events_and_gtis(fits_file, additional_columns=None,
                         gtistring='GTI,STDGTI',
                         gti_file=None, hduname='EVENTS', column='TIME'):

    """Load event lists and GTIs from one or more files.

    Loads event list from HDU EVENTS of file fits_file, with Good Time
    intervals. Optionally, returns additional columns of data from the same
    HDU of the events.

    Parameters
    ----------
    fits_file : str
    return_limits: bool, optional
        Return the TSTART and TSTOP keyword values
    additional_columns: list of str, optional
        A list of keys corresponding to the additional columns to extract from
        the event HDU (ex.: ['PI', 'X'])

    Returns
    -------
    ev_list : array-like
    gtis: [[gti0_0, gti0_1], [gti1_0, gti1_1], ...]
    additional_data: dict
        A dictionary, where each key is the one specified in additional_colums.
        The data are an array with the values of the specified column in the
        fits file.
    t_start : float
    t_stop : float
    """

    gtistring = assign_value_if_none(gtistring, 'GTI,STDGTI')
    lchdulist = fits.open(fits_file)

    # Load data table
    try:
        lctable = lchdulist[hduname].data
    except:  # pragma: no cover
        logging.warning('HDU %s not found. Trying first extension' % hduname)
        lctable = lchdulist[1].data

    # Read event list
    ev_list = np.array(lctable.field(column), dtype=np.longdouble)

    # Read TIMEZERO keyword and apply it to events
    try:
        timezero = np.longdouble(lchdulist[1].header['TIMEZERO'])
    except:  # pragma: no cover
        logging.warning("No TIMEZERO in file")
        timezero = np.longdouble(0.)

    ev_list += timezero

    # Read TSTART, TSTOP from header
    try:
        t_start = np.longdouble(lchdulist[1].header['TSTART'])
        t_stop = np.longdouble(lchdulist[1].header['TSTOP'])
    except:  # pragma: no cover
        logging.warning("Tstart and Tstop error. using defaults")
        t_start = ev_list[0]
        t_stop = ev_list[-1]

    # Read and handle GTI extension
    accepted_gtistrings = gtistring.split(',')

    if gti_file is None:
        # Select first GTI with accepted name
        try:
            gti_list = \
                _get_gti_from_extension(
                    lchdulist, accepted_gtistrings=accepted_gtistrings)
        except:  # pragma: no cover
            warnings.warn("No extensions found with a valid name. "
                          "Please check the `accepted_gtistrings` values.")
            gti_list = np.array([[t_start, t_stop]],
                                dtype=np.longdouble)
    else:
        gti_list = load_gtis(gti_file, gtistring)

    additional_data = _get_additional_data(lctable, additional_columns)

    lchdulist.close()

    # Sort event list
    order = np.argsort(ev_list)
    ev_list = ev_list[order]

    additional_data = order_list_of_arrays(additional_data, order)

    returns = _empty()
    returns.ev_list = ev_list
    returns.gti_list = gti_list
    returns.additional_data = additional_data
    returns.t_start = t_start
    returns.t_stop = t_stop

    return returns

class _empty():
    def __init__(self):
        pass

def mkdir_p(path):  # pragma: no cover
    """Safe mkdir function.

    Parameters
    ----------
    path : str
        Name of the directory/ies to create

    Notes
    -----
    Found at
    http://stackoverflow.com/questions/600268/mkdir-p-functionality-in-python
    """
    import os
    import errno
    try:
        os.makedirs(path)
    except OSError as exc:  # Python >2.5
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise

def read_header_key(fits_file, key, hdu=1):
    """Read the header key key from HDU hdu of the file fits_file.

    Parameters
    ----------
    fits_file: str
    key: str
        The keyword to be read

    Other Parameters
    ----------------
    hdu : int
    """

    hdulist = fits.open(fits_file)
    try:
        value = hdulist[hdu].header[key]
    except:  # pragma: no cover
        value = ''
    hdulist.close()
    return value

def ref_mjd(fits_file, hdu=1):
    """Read MJDREFF+ MJDREFI or, if failed, MJDREF, from the FITS header.

    Parameters
    ----------
    fits_file : str

    Returns
    -------
    mjdref : numpy.longdouble
        the reference MJD

    Other Parameters
    ----------------
    hdu : int
    """
    import collections

    if isinstance(fits_file, collections.Iterable) and\
            not is_string(fits_file):  # pragma: no cover
        fits_file = fits_file[0]
        logging.info("opening %s" % fits_file)

    hdulist = fits.open(fits_file)

    ref_mjd_val = high_precision_keyword_read(hdulist[hdu].header, "MJDREF")

    hdulist.close()
    return ref_mjd_val

def common_name(str1, str2, default='common'):
    """Strip two strings of the letters not in common.

    Filenames must be of same length and only differ by a few letters.

    Parameters
    ----------
    str1 : str
    str2 : str

    Returns
    -------
    common_str : str
        A string containing the parts of the two names in common

    Other Parameters
    ----------------
    default : str
        The string to return if common_str is empty
    """
    if not len(str1) == len(str2):
        return default
    common_str = ''
    # Extract the MP root of the name (in case they're event files)

    for i, letter in enumerate(str1):
        if str2[i] == letter:
            common_str += letter
    # Remove leading and trailing underscores and dashes
    common_str = common_str.rstrip('_').rstrip('-')
    common_str = common_str.lstrip('_').lstrip('-')
    if common_str == '':
        common_str = default
    logging.debug('common_name: %s %s -> %s' % (str1, str2, common_str))
    return common_str

def _save_pickle_object(object, filename, **kwargs):

    if 'save_as_dict' in locals():
        # Get all object's attributes and its values in a dictionary format
        items = vars(object)
        with open(filename, "wb" ) as f:
            pickle.dump(items, f)

    else:
        with open(filename, "wb" ) as f:
            pickle.dump(object, f)

def _retrieve_pickle_object(filename):
    with open(filename, "rb" ) as f:
        return pickle.load(f)

def _save_hdf5_object(object, filename):
    pass

def _retrieve_hdf5_object(filename, **kwargs):
    pass

def _save_ascii_object(object, filename, **kwargs):
    """
    Save an array to a text file.

    Parameters
    ----------
    object : numpy.ndarray
        An array with the data to be saved

    filename : str
        The file name to save to

    Other Parameters
    ----------------
    kwargs : any keyword argument taken by `numpy.loadtxt`

    """
    np.savetxt(object, filename, **kwargs)
    pass

def _retrieve_ascii_object(filename, **kwargs):
    """
    Helper function to retrieve ascii objects from file.
    Uses astropy.Table for reading and storing the data.

    Parameters
    ----------
    filename : str
        The name of the file with the data to be retrieved.

    Additional Keyword Parameters
    -----------------------------
    usecols : {int | iterable}
        The indices of the columns in the file to be returned.
        By default, all columns will be returned

    skiprows : int
        The number of rows at the beginning to skip
        By default, no rows will be skipped.

    names : iterable
        A list of column names to be attached to the columns.
        By default, no column names are added, unless they are specified
        in the file header and can be read by astropy.Table.read
        automatically.

    Returns
    -------
    data : astropy.Table object
        An astropy.Table object with the data from the file

    TODO: Add example
    """
    assert isinstance(filename, str), "filename must be string!"

    if 'usecols' in list(kwargs.keys()):
        assert np.size(kwargs['usecols']) == 2, "Need to define two columns"
        usecols = kwargs["usecols"]
    else:
        usecols = None

    if 'skiprows' in list(kwargs.keys()):
        assert isinstance(kwargs["skiprows"], int)
        skiprows = kwargs["skiprows"]
    else:
        skiprows = 0

    if "names" in list(kwargs.keys()):
        names = kwargs["names"]
    else:
        names = None

    data = Table.read(filename, data_start=skiprows,
                      names=names)

    if usecols is None:
        return data
    else:
        colnames = np.array(data.colnames)
        cols = colnames[usecols]

        return data[cols]

def write(input_, filename, format_='pickle', **kwargs):
    """
    Pickle a class instance. For parameters depending on
    `format_`, see individual function definitions.

    Parameters
    ----------
    object: a class instance
    filename: str
        name of the file to be created.
    format_: str
        pickle, hdf5, ascii ...

    save_as_dict: boolean
        For compatibility with MaLTpyNT, save_as_dict should be true. Set
        it to 'False' if intention is to store input as class object.
    """

    if format_ == 'pickle':
        _save_pickle_object(input_, filename, **kwargs)

    elif format_ == 'hdf5':
        _save_hdf5_object(input_, filename)

    elif format_ == 'ascii':
        _save_ascii_object(input_, filename)

def read(filename, format_='pickle', **kwargs):
    """
    Return a pickled class instance.

    Parameters
    ----------
    filename: str
        name of the file to be retrieved.
    format_: str
        pickle, hdf5, ascii ...
    """

    if format_ == 'pickle':
        return _retrieve_pickle_object(filename)

    elif format_ == 'hdf5':
        return _retrieve_hdf5_object(filename, **kwargs)

    elif format_ == 'ascii':
        return _retrieve_ascii_object(filename, **kwargs)

def savefig(filename, **kwargs):
    """
    Save a figure plotted by Matplotlib.

    Note : This function is supposed to be used after the ``plot``
    function. Otherwise it will save a blank image with no plot.

    Parameters
    ----------
    filename : str
        The name of the image file. Extension must be specified in the
        file name. For example filename with `.png` extension will give a
        rasterized image while `.pdf` extension will give a vectorized
        output.

    kwargs : keyword arguments
        Keyword arguments to be passed to ``savefig`` function of
        ``matplotlib.pyplot``. For example use `bbox_inches='tight'` to
        remove the undesirable whitepace around the image.
    """

    try:
        import matplotlib.pyplot as plt
    except ImportError:
        raise ImportError("Matplotlib required for savefig()")

    if not plt.fignum_exists(1):
        utils.simon("use ``plot`` function to plot the image first and "
                    "then use ``savefig`` to save the figure.")

    plt.savefig(filename, **kwargs)
<commit_msg>Fixed bug in saving ascii in io.py<commit_after>from __future__ import (absolute_import, unicode_literals, division,
                        print_function)
import numpy as np
import logging
import warnings
import os

from astropy.io import fits
from astropy.table import Table

import stingray.utils as utils
from .utils import order_list_of_arrays, is_string
from .utils import assign_value_if_none

from .gti import  _get_gti_from_extension, load_gtis

try:
    # Python 2
    import cPickle as pickle
except:
    # Python 3
    import pickle

def get_file_extension(fname):
    """Get the extension from the file name."""
    return os.path.splitext(fname)[1]

def high_precision_keyword_read(hdr, keyword):
    """Read FITS header keywords, also if split in two.

    In the case where the keyword is split in two, like

        MJDREF = MJDREFI + MJDREFF

    in some missions, this function returns the summed value. Otherwise, the
    content of the single keyword

    Parameters
    ----------
    hdr : dict_like
        The FITS header structure, or a dictionary
    keyword : str
        The key to read in the header

    Returns
    -------
    value : long double
        The value of the key, or None if something went wrong

    """
    try:
        value = np.longdouble(hdr[keyword])
        return value
    except:
        pass
    try:
        if len(keyword) == 8:
            keyword = keyword[:7]
        value = np.longdouble(hdr[keyword + 'I'])
        value += np.longdouble(hdr[keyword + 'F'])
        return value
    except:
        return None

def _get_additional_data(lctable, additional_columns):
    additional_data = {}
    if additional_columns is not None:
        for a in additional_columns:
            try:
                additional_data[a] = np.array(lctable.field(a))
            except:  # pragma: no cover
                if a == 'PI':
                    logging.warning('Column PI not found. Trying with PHA')
                    additional_data[a] = np.array(lctable.field('PHA'))
                else:
                    raise Exception('Column' + a + 'not found')

    return additional_data

def load_events_and_gtis(fits_file, additional_columns=None,
                         gtistring='GTI,STDGTI',
                         gti_file=None, hduname='EVENTS', column='TIME'):

    """Load event lists and GTIs from one or more files.

    Loads event list from HDU EVENTS of file fits_file, with Good Time
    intervals. Optionally, returns additional columns of data from the same
    HDU of the events.

    Parameters
    ----------
    fits_file : str
    return_limits: bool, optional
        Return the TSTART and TSTOP keyword values
    additional_columns: list of str, optional
        A list of keys corresponding to the additional columns to extract from
        the event HDU (ex.: ['PI', 'X'])

    Returns
    -------
    ev_list : array-like
    gtis: [[gti0_0, gti0_1], [gti1_0, gti1_1], ...]
    additional_data: dict
        A dictionary, where each key is the one specified in additional_colums.
        The data are an array with the values of the specified column in the
        fits file.
    t_start : float
    t_stop : float
    """

    gtistring = assign_value_if_none(gtistring, 'GTI,STDGTI')
    lchdulist = fits.open(fits_file)

    # Load data table
    try:
        lctable = lchdulist[hduname].data
    except:  # pragma: no cover
        logging.warning('HDU %s not found. Trying first extension' % hduname)
        lctable = lchdulist[1].data

    # Read event list
    ev_list = np.array(lctable.field(column), dtype=np.longdouble)

    # Read TIMEZERO keyword and apply it to events
    try:
        timezero = np.longdouble(lchdulist[1].header['TIMEZERO'])
    except:  # pragma: no cover
        logging.warning("No TIMEZERO in file")
        timezero = np.longdouble(0.)

    ev_list += timezero

    # Read TSTART, TSTOP from header
    try:
        t_start = np.longdouble(lchdulist[1].header['TSTART'])
        t_stop = np.longdouble(lchdulist[1].header['TSTOP'])
    except:  # pragma: no cover
        logging.warning("Tstart and Tstop error. using defaults")
        t_start = ev_list[0]
        t_stop = ev_list[-1]

    # Read and handle GTI extension
    accepted_gtistrings = gtistring.split(',')

    if gti_file is None:
        # Select first GTI with accepted name
        try:
            gti_list = \
                _get_gti_from_extension(
                    lchdulist, accepted_gtistrings=accepted_gtistrings)
        except:  # pragma: no cover
            warnings.warn("No extensions found with a valid name. "
                          "Please check the `accepted_gtistrings` values.")
            gti_list = np.array([[t_start, t_stop]],
                                dtype=np.longdouble)
    else:
        gti_list = load_gtis(gti_file, gtistring)

    additional_data = _get_additional_data(lctable, additional_columns)

    lchdulist.close()

    # Sort event list
    order = np.argsort(ev_list)
    ev_list = ev_list[order]

    additional_data = order_list_of_arrays(additional_data, order)

    returns = _empty()
    returns.ev_list = ev_list
    returns.gti_list = gti_list
    returns.additional_data = additional_data
    returns.t_start = t_start
    returns.t_stop = t_stop

    return returns

class _empty():
    def __init__(self):
        pass

def mkdir_p(path):  # pragma: no cover
    """Safe mkdir function.

    Parameters
    ----------
    path : str
        Name of the directory/ies to create

    Notes
    -----
    Found at
    http://stackoverflow.com/questions/600268/mkdir-p-functionality-in-python
    """
    import os
    import errno
    try:
        os.makedirs(path)
    except OSError as exc:  # Python >2.5
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise

def read_header_key(fits_file, key, hdu=1):
    """Read the header key key from HDU hdu of the file fits_file.

    Parameters
    ----------
    fits_file: str
    key: str
        The keyword to be read

    Other Parameters
    ----------------
    hdu : int
    """

    hdulist = fits.open(fits_file)
    try:
        value = hdulist[hdu].header[key]
    except:  # pragma: no cover
        value = ''
    hdulist.close()
    return value

def ref_mjd(fits_file, hdu=1):
    """Read MJDREFF+ MJDREFI or, if failed, MJDREF, from the FITS header.

    Parameters
    ----------
    fits_file : str

    Returns
    -------
    mjdref : numpy.longdouble
        the reference MJD

    Other Parameters
    ----------------
    hdu : int
    """
    import collections

    if isinstance(fits_file, collections.Iterable) and\
            not is_string(fits_file):  # pragma: no cover
        fits_file = fits_file[0]
        logging.info("opening %s" % fits_file)

    hdulist = fits.open(fits_file)

    ref_mjd_val = high_precision_keyword_read(hdulist[hdu].header, "MJDREF")

    hdulist.close()
    return ref_mjd_val

def common_name(str1, str2, default='common'):
    """Strip two strings of the letters not in common.

    Filenames must be of same length and only differ by a few letters.

    Parameters
    ----------
    str1 : str
    str2 : str

    Returns
    -------
    common_str : str
        A string containing the parts of the two names in common

    Other Parameters
    ----------------
    default : str
        The string to return if common_str is empty
    """
    if not len(str1) == len(str2):
        return default
    common_str = ''
    # Extract the MP root of the name (in case they're event files)

    for i, letter in enumerate(str1):
        if str2[i] == letter:
            common_str += letter
    # Remove leading and trailing underscores and dashes
    common_str = common_str.rstrip('_').rstrip('-')
    common_str = common_str.lstrip('_').lstrip('-')
    if common_str == '':
        common_str = default
    logging.debug('common_name: %s %s -> %s' % (str1, str2, common_str))
    return common_str

def _save_pickle_object(object, filename, **kwargs):

    if 'save_as_dict' in locals():
        # Get all object's attributes and its values in a dictionary format
        items = vars(object)
        with open(filename, "wb" ) as f:
            pickle.dump(items, f)

    else:
        with open(filename, "wb" ) as f:
            pickle.dump(object, f)

def _retrieve_pickle_object(filename):
    with open(filename, "rb" ) as f:
        return pickle.load(f)

def _save_hdf5_object(object, filename):
    pass

def _retrieve_hdf5_object(filename, **kwargs):
    pass

def _save_ascii_object(object, filename, **kwargs):
    """
    Save an array to a text file.

    Parameters
    ----------
    object : numpy.ndarray
        An array with the data to be saved

    filename : str
        The file name to save to

    Other Parameters
    ----------------
    kwargs : any keyword argument taken by `numpy.savetxt`

    """
    np.savetxt(filename, object, **kwargs)
    pass

def _retrieve_ascii_object(filename, **kwargs):
    """
    Helper function to retrieve ascii objects from file.
    Uses astropy.Table for reading and storing the data.

    Parameters
    ----------
    filename : str
        The name of the file with the data to be retrieved.

    Additional Keyword Parameters
    -----------------------------
    usecols : {int | iterable}
        The indices of the columns in the file to be returned.
        By default, all columns will be returned

    skiprows : int
        The number of rows at the beginning to skip
        By default, no rows will be skipped.

    names : iterable
        A list of column names to be attached to the columns.
        By default, no column names are added, unless they are specified
        in the file header and can be read by astropy.Table.read
        automatically.

    Returns
    -------
    data : astropy.Table object
        An astropy.Table object with the data from the file

    TODO: Add example
    """
    assert isinstance(filename, str), "filename must be string!"

    if 'usecols' in list(kwargs.keys()):
        assert np.size(kwargs['usecols']) == 2, "Need to define two columns"
        usecols = kwargs["usecols"]
    else:
        usecols = None

    if 'skiprows' in list(kwargs.keys()):
        assert isinstance(kwargs["skiprows"], int)
        skiprows = kwargs["skiprows"]
    else:
        skiprows = 0

    if "names" in list(kwargs.keys()):
        names = kwargs["names"]
    else:
        names = None

    data = Table.read(filename, data_start=skiprows,
                      names=names)

    if usecols is None:
        return data
    else:
        colnames = np.array(data.colnames)
        cols = colnames[usecols]

        return data[cols]

def write(input_, filename, format_='pickle', **kwargs):
    """
    Pickle a class instance. For parameters depending on
    `format_`, see individual function definitions.

    Parameters
    ----------
    object: a class instance
    filename: str
        name of the file to be created.
    format_: str
        pickle, hdf5, ascii ...

    save_as_dict: boolean
        For compatibility with MaLTpyNT, save_as_dict should be true. Set
        it to 'False' if intention is to store input as class object.
    """

    if format_ == 'pickle':
        _save_pickle_object(input_, filename, **kwargs)

    elif format_ == 'hdf5':
        _save_hdf5_object(input_, filename)

    elif format_ == 'ascii':
        _save_ascii_object(input_, filename)

def read(filename, format_='pickle', **kwargs):
    """
    Return a pickled class instance.

    Parameters
    ----------
    filename: str
        name of the file to be retrieved.
    format_: str
        pickle, hdf5, ascii ...
    """

    if format_ == 'pickle':
        return _retrieve_pickle_object(filename)

    elif format_ == 'hdf5':
        return _retrieve_hdf5_object(filename, **kwargs)

    elif format_ == 'ascii':
        return _retrieve_ascii_object(filename, **kwargs)

def savefig(filename, **kwargs):
    """
    Save a figure plotted by Matplotlib.

    Note : This function is supposed to be used after the ``plot``
    function. Otherwise it will save a blank image with no plot.

    Parameters
    ----------
    filename : str
        The name of the image file. Extension must be specified in the
        file name. For example filename with `.png` extension will give a
        rasterized image while `.pdf` extension will give a vectorized
        output.

    kwargs : keyword arguments
        Keyword arguments to be passed to ``savefig`` function of
        ``matplotlib.pyplot``. For example use `bbox_inches='tight'` to
        remove the undesirable whitepace around the image.
    """

    try:
        import matplotlib.pyplot as plt
    except ImportError:
        raise ImportError("Matplotlib required for savefig()")

    if not plt.fignum_exists(1):
        utils.simon("use ``plot`` function to plot the image first and "
                    "then use ``savefig`` to save the figure.")

    plt.savefig(filename, **kwargs)
<|endoftext|>


<commit_before># author: G. Alomar
from cassandra.cluster import Cluster
from hecuba.datastore import *
from hecuba.iter import Block
from hecuba.dict import *
from hecuba.storageobj import *
from hecuba.storageobjix import *
import collections

from conf.hecuba_params import *
from conf.apppath import apppath

def start_task(params):
    if not 'prefetch_activated' in globals():
        global prefetch_activated
        prefetch_activated = True
    if not 'batch_activated' in globals():
        global batch_activated
        batch_activated = True
    if prefetch_activated or batch_activated:
        path = apppath + '/conf/imports.py'
        file = open(path, 'r')
        for line in file:
            exec line

    if batch_activated:
        for param in params:
            if isinstance(param, StorageObj) or isinstance(param, StorageObjIx) or isinstance(param, Block):
                param.cntxt = context(param)
                param.cntxt.__enter__()

def end_task(params):
    if not 'prefetch_activated' in globals():
        global prefetch_activated
        prefetch_activated = True
    if not 'batch_activated' in globals():
        global batch_activated
        batch_activated = True
    if prefetch_activated or batch_activated:
        path = apppath + '/conf/imports.py'
        file = open(path, 'r')
        for line in file:
            exec line

    if batch_activated:
        for param in params:
            if isinstance(param, Block) or isinstance(param, StorageObj) or isinstance(param, StorageObjIx):
                try:
                    param.cntxt.__exit__()
                except Exception as e:
                    print "error trying to exit context:", e

    if prefetch_activated:
        for param in params:
            if isinstance(param, Block):
                keys = param.storageobj.keyList[param.storageobj.__class__.__name__]
                exec("persistentdict = param.storageobj." + str(keys[0]))
                if persistentdict.prefetch == True:
                    try:
                        persistentdict.end_prefetch()
                    except Exception as e:
                        print "error trying to prefetch:", e

    if not 'statistics_activated' in globals():
        global statistics_activated
        statistics_activated = True
    if statistics_activated:
        for param in params:
            if isinstance(param, Block):
                param.storageobj.statistics()
            if isinstance(param, StorageObj) or isinstance(param, StorageObjIx):
                param.statistics()

def getByID(objid):
    path = apppath + '/conf/imports.py'
    file = open(path, 'r')
    for line in file:
        exec(line)
    objidsplit = objid.split("_")
    print "objidsplit:", objidsplit

    if len(objidsplit) == 2:
        result = ''
        exec('result = ' + objidsplit[0] + '()')
        print "str(objidsplit[0]):", str(objidsplit[0])
        result.getByName(str(objidsplit[0]))
        return result

    else:
        cluster = Cluster(contact_points=contact_names, port=nodePort, protocol_version=2)
        session = cluster.connect()
        (ksp,tab,dict_name,obj_type,entryPoint,port,tokens) = session.execute("SELECT ksp,tab,dict_name,obj_type,entrypoint,port,tkns FROM hecuba.blocks WHERE blockid = %s",(objid))[0]
        if obj_type == 'qbeast':
            print "indexed storageobj"
            blockid = objid
            metadata = cluster.metadata
            tokenmap = metadata.token_map.token_to_host_owner
            odtokenmap = collections.OrderedDict(sorted(tokenmap.items()))
            if not 'prefetch_activated' in globals():
                global prefetch_activated
                prefetch_activated = True
            for position in tokens:
                for key, val in odtokenmap.iteritems():
                    # (self, peer,        keynames,           tablename,     blockkeyspace, myuuid)
                    b = IxBlock(str(val), str(objidsplit[1]), tab,           ksp,           objid)
                    exec("b.storageobj = " + str(dict_name) + "('" + str(dict_name) + "')")
                    if prefetch_activated:
                        b.storageobj.init_prefetch(b)
                    session.shutdown()
                    cluster.shutdown()
                    return b
        else:
            print "normal storageobj"
            # blockKeyspace = objidsplit[0]
            # blockKeyNames = str(objidsplit[1])
            # blockTableName = objidsplit[2]
            # blockRange = objidsplit[3]
            blockranges = objidsplit[3:len(objidsplit)]
            blockrangesf = ''
            for ind, val in enumerate(blockranges):
                if ind < (len(blockranges)-1):
                    blockrangesf += str(val) + '_'
                else:
                    blockrangesf += str(val)
            cluster = Cluster(contact_points=contact_names, port=nodePort, protocol_version=2)
            session = cluster.connect()
            metadata = cluster.metadata
            tokenmap = metadata.token_map.token_to_host_owner
            odtokenmap = collections.OrderedDict(sorted(tokenmap.items()))
            position = 0
            if not 'prefetch_activated' in globals():
                global prefetch_activated
            prefetch_activated = True
            for key, val in odtokenmap.iteritems():
                if str(position) == objidsplit[3]:
                    b = Block((str(val), blockrangesf), str(objidsplit[1]), objidsplit[2], objidsplit[0])
                    exec("b.storageobj = " + str(objidsplit[2]) + "('" + str(objidsplit[2]) + "')")
                    if prefetch_activated:
                        b.storageobj.init_prefetch(b)
                    session.shutdown()
                    cluster.shutdown()
                    return b
                position += 1
<commit_msg>fixing those errors...<commit_after># author: G. Alomar
from cassandra.cluster import Cluster
from hecuba.datastore import *
from hecuba.iter import Block
from hecuba.dict import *
from hecuba.storageobj import *
from hecuba.storageobjix import *
import collections

from conf.hecuba_params import *
from conf.apppath import apppath

def start_task(params):
    if not 'prefetch_activated' in globals():
        global prefetch_activated
        prefetch_activated = True
    if not 'batch_activated' in globals():
        global batch_activated
        batch_activated = True
    if prefetch_activated or batch_activated:
        path = apppath + '/conf/imports.py'
        file = open(path, 'r')
        for line in file:
            exec line

    if batch_activated:
        for param in params:
            if isinstance(param, StorageObj) or isinstance(param, StorageObjIx) or isinstance(param, Block):
                param.cntxt = context(param)
                param.cntxt.__enter__()

def end_task(params):
    if not 'prefetch_activated' in globals():
        global prefetch_activated
        prefetch_activated = True
    if not 'batch_activated' in globals():
        global batch_activated
        batch_activated = True
    if prefetch_activated or batch_activated:
        path = apppath + '/conf/imports.py'
        file = open(path, 'r')
        for line in file:
            exec line

    if batch_activated:
        for param in params:
            if isinstance(param, Block) or isinstance(param, StorageObj) or isinstance(param, StorageObjIx):
                try:
                    param.cntxt.__exit__()
                except Exception as e:
                    print "error trying to exit context:", e

    if prefetch_activated:
        for param in params:
            if isinstance(param, Block):
                keys = param.storageobj.keyList[param.storageobj.__class__.__name__]
                exec("persistentdict = param.storageobj." + str(keys[0]))
                if persistentdict.prefetch == True:
                    try:
                        persistentdict.end_prefetch()
                    except Exception as e:
                        print "error trying to prefetch:", e

    if not 'statistics_activated' in globals():
        global statistics_activated
        statistics_activated = True
    if statistics_activated:
        for param in params:
            if isinstance(param, Block):
                param.storageobj.statistics()
            if isinstance(param, StorageObj) or isinstance(param, StorageObjIx):
                param.statistics()

def getByID(objid):
    path = apppath + '/conf/imports.py'
    file = open(path, 'r')
    for line in file:
        exec(line)
    objidsplit = objid.split("_")
    print "objidsplit:", objidsplit

    if len(objidsplit) == 2:
        result = ''
        exec('result = ' + objidsplit[0] + '()')
        print "str(objidsplit[0]):", str(objidsplit[0])
        result.getByName(str(objidsplit[0]))
        return result

    else:
        cluster = Cluster(contact_points=contact_names, port=nodePort, protocol_version=2)
        session = cluster.connect()
        #(ksp,tab,dict_name,obj_type,entryPoint,port,tokens) = session.execute("SELECT ksp,tab,dict_name,obj_type,entrypoint,port,tkns FROM hecuba.blocks WHERE blockid = %s",(objid))[0]
        try:
            ksp =        session.execute("SELECT ksp FROM hecuba.blocks WHERE blockid = %s",(objid))[0]
        except Exception as e:
            print "Error importing values in getByID:", e
        try:
            tab =        session.execute("SELECT tab FROM hecuba.blocks WHERE blockid = %s",(objid))[0]
        except Exception as e:
            print "Error importing values in getByID:", e
        try:
            dict_name =  session.execute("SELECT dict_name FROM hecuba.blocks WHERE blockid = %s",(objid))[0]
        except Exception as e:
            print "Error importing values in getByID:", e
        try:
            obj_type =   session.execute("SELECT obj_type FROM hecuba.blocks WHERE blockid = %s",(objid))[0]
        except Exception as e:
            print "Error importing values in getByID:", e
        try:
            entryPoint = session.execute("SELECT entrypoint FROM hecuba.blocks WHERE blockid = %s",(objid))[0]
        except Exception as e:
            print "Error importing values in getByID:", e
        try:
            port =       session.execute("SELECT port FROM hecuba.blocks WHERE blockid = %s",(objid))[0]
        except Exception as e:
            print "Error importing values in getByID:", e
        try:
            tokens =     session.execute("SELECT tkns FROM hecuba.blocks WHERE blockid = %s",(objid))[0]
        except Exception as e:
            print "Error importing values in getByID:", e
        if obj_type == 'qbeast':
            print "indexed storageobj"
            blockid = objid
            metadata = cluster.metadata
            tokenmap = metadata.token_map.token_to_host_owner
            odtokenmap = collections.OrderedDict(sorted(tokenmap.items()))
            if not 'prefetch_activated' in globals():
                global prefetch_activated
                prefetch_activated = True
            for position in tokens:
                for key, val in odtokenmap.iteritems():
                    # (self, peer,        keynames,           tablename,     blockkeyspace, myuuid)
                    b = IxBlock(str(val), str(objidsplit[1]), tab,           ksp,           objid)
                    exec("b.storageobj = " + str(dict_name) + "('" + str(dict_name) + "')")
                    if prefetch_activated:
                        b.storageobj.init_prefetch(b)
                    session.shutdown()
                    cluster.shutdown()
                    return b
        else:
            print "normal storageobj"
            # blockKeyspace = objidsplit[0]
            # blockKeyNames = str(objidsplit[1])
            # blockTableName = objidsplit[2]
            # blockRange = objidsplit[3]
            blockranges = objidsplit[3:len(objidsplit)]
            blockrangesf = ''
            for ind, val in enumerate(blockranges):
                if ind < (len(blockranges)-1):
                    blockrangesf += str(val) + '_'
                else:
                    blockrangesf += str(val)
            cluster = Cluster(contact_points=contact_names, port=nodePort, protocol_version=2)
            session = cluster.connect()
            metadata = cluster.metadata
            tokenmap = metadata.token_map.token_to_host_owner
            odtokenmap = collections.OrderedDict(sorted(tokenmap.items()))
            position = 0
            if not 'prefetch_activated' in globals():
                global prefetch_activated
            prefetch_activated = True
            for key, val in odtokenmap.iteritems():
                if str(position) == objidsplit[3]:
                    b = Block((str(val), blockrangesf), str(objidsplit[1]), objidsplit[2], objidsplit[0])
                    exec("b.storageobj = " + str(objidsplit[2]) + "('" + str(objidsplit[2]) + "')")
                    if prefetch_activated:
                        b.storageobj.init_prefetch(b)
                    session.shutdown()
                    cluster.shutdown()
                    return b
                position += 1
<|endoftext|>


<commit_before>import sys
from functools import wraps

from django.http import Http404
from django.conf import settings
from django.utils.importlib import import_module
from django.db.models.signals import pre_save, post_save
from django.dispatch import receiver

from .models import Page, pre_route, post_route

def pre_url(route_url):
    def outer(f):
        @wraps(f)
        def wrapper(sender, url, **kwargs):
            if route_url == url:
                return f(request=sender, url=url)
            else:
                return False
        pre_route.connect(wrapper)
        return wrapper
    return outer

def get_page_from_url(url):
    if url[-1] != '/':
        url = url + '/'

    if url[0] != '/':
        url = '/' + url
    try:
        return Page.objects.select_related().live().get(url=url)
    except Page.DoesNotExist:
        raise Http404

<commit_msg>default images<commit_after>import sys
from functools import wraps

from django.http import Http404
from django.conf import settings
from django.utils.importlib import import_module
from django.db.models.signals import pre_save, post_save
from django.dispatch import receiver

from .models import Page, pre_route, post_route, Image

def pre_url(route_url):
    def outer(f):
        @wraps(f)
        def wrapper(sender, url, **kwargs):
            if route_url == url:
                return f(request=sender, url=url)
            else:
                return False
        pre_route.connect(wrapper)
        return wrapper
    return outer

def get_page_from_url(url):
    if url[-1] != '/':
        url = url + '/'

    if url[0] != '/':
        url = '/' + url
    try:
        return Page.objects.select_related().live().get(url=url)
    except Page.DoesNotExist:
        raise Http404

def get_default_image():
    try:
        return Image(image=settings.DEFAULT_IMAGE)
    except AttributeError:
        return None<|endoftext|>


<commit_before>from abusehelper.core import bot, config, startup, runtime

def iter_dots(iterable):
    for obj in iterable:
        dot = getattr(obj, "__dot__", None)
        if callable(dot):
            yield dot()

class Dot(object):
    def __init__(self, *configs):
        self.configs = tuple(config.flatten(configs))

    def __dot__(self):
        return self

    def services(self):
        services = set()
        for dot in iter_dots(self.configs):
            services.update(dot.services())

        for conf in startup.iter_startups(self.configs):
            service = dict(startup.params).get("bot_name", None)
            services.add(service)

        services.discard(None)
        return services

    def sessions(self):
        sessions = set()
        for dot in iter_dots(self.configs):
            sessions.update(dot.sessions())

        for session in runtime.iter_runtimes(self.configs):
            sessions.add(session)

        return sessions

class DotBot(bot.Bot):
    bot_name = None

    config = bot.Param("configuration module")
    show_startups = bot.BoolParam()
    show_attributes = bot.BoolParam()

    def run(self):
        print "digraph G {"
        print "node [ shape=box, style=filled, color=lightgrey ];"

        dot = Dot(config.load_configs(self.config))
        services = dot.services()
        sessions = dot.sessions()

        if self.show_startups:
            missing = lambda x: x not in services
        else:
            missing = lambda x: False

        for session in sessions:
            conf = dict(session.conf)
            path = session.path
            src = conf.pop("src_room", None)
            dst = conf.pop("dst_room", None)

            node = None
            if session.service != "roomgraph":
                node = "node " + session.service + " " + (src or "@") + " " + (dst or "@")
                if missing(session.service):
                    print '"%s" [ label="%s", shape=ellipse, fontsize=12, fontcolor=white, color=red ];' % (node, session.service)
                else:
                    print '"%s" [ label="%s", shape=ellipse, fontsize=12, style="" ];' % (node, session.service)

            name = session.service

            label_list = list()
            if path:
                label_list.append(".".join(path))
            if self.show_attributes:
                for item in conf.items():
                    label_list.append("  %r=%r" % item)
            label = "\\n".join(label_list)

            color = "red" if missing(session.service) else "black"
            if src is None:
                if node is not None and dst is not None:
                    print '"%s"->"%s" [ label="%s", fontsize=10, color="%s" ];' % (node, dst, label, color)
            else:
                if node is not None and dst is None:
                    print '"%s"->"%s" [ label="%s", fontsize=10, color="%s" ];' % (src, node, label, color)
                elif node is None and dst is not None:
                    print '"%s"->"%s" [ label="%s", fontsize=10, color="%s" ];' % (src, dst, label, color)
                elif node is not None and dst is not None:
                    print '"%s"->"%s" [ label="%s", fontsize=10, color="%s" ];' % (src, node, "", color)
                    print '"%s"->"%s" [ label="%s", fontsize=10, color="%s" ];' % (node, dst, label, color)

        print '}'

if __name__ == "__main__":
    DotBot.from_command_line().execute()
<commit_msg>contrib/runtime2dot: Fix a wrong variable name.<commit_after>from abusehelper.core import bot, config, startup, runtime

def iter_dots(iterable):
    for obj in iterable:
        dot = getattr(obj, "__dot__", None)
        if callable(dot):
            yield dot()

class Dot(object):
    def __init__(self, *configs):
        self.configs = tuple(config.flatten(configs))

    def __dot__(self):
        return self

    def services(self):
        services = set()
        for dot in iter_dots(self.configs):
            services.update(dot.services())

        for conf in startup.iter_startups(self.configs):
            service = dict(conf.params).get("bot_name", None)
            services.add(service)

        services.discard(None)
        return services

    def sessions(self):
        sessions = set()
        for dot in iter_dots(self.configs):
            sessions.update(dot.sessions())

        for session in runtime.iter_runtimes(self.configs):
            sessions.add(session)

        return sessions

class DotBot(bot.Bot):
    bot_name = None

    config = bot.Param("configuration module")
    show_startups = bot.BoolParam()
    show_attributes = bot.BoolParam()

    def run(self):
        print "digraph G {"
        print "node [ shape=box, style=filled, color=lightgrey ];"

        dot = Dot(config.load_configs(self.config))
        services = dot.services()
        sessions = dot.sessions()

        if self.show_startups:
            missing = lambda x: x not in services
        else:
            missing = lambda x: False

        for session in sessions:
            conf = dict(session.conf)
            path = session.path
            src = conf.pop("src_room", None)
            dst = conf.pop("dst_room", None)

            node = None
            if session.service != "roomgraph":
                node = "node " + session.service + " " + (src or "@") + " " + (dst or "@")
                if missing(session.service):
                    print '"%s" [ label="%s", shape=ellipse, fontsize=12, fontcolor=white, color=red ];' % (node, session.service)
                else:
                    print '"%s" [ label="%s", shape=ellipse, fontsize=12, style="" ];' % (node, session.service)

            name = session.service

            label_list = list()
            if path:
                label_list.append(".".join(path))
            if self.show_attributes:
                for item in conf.items():
                    label_list.append("  %r=%r" % item)
            label = "\\n".join(label_list)

            color = "red" if missing(session.service) else "black"
            if src is None:
                if node is not None and dst is not None:
                    print '"%s"->"%s" [ label="%s", fontsize=10, color="%s" ];' % (node, dst, label, color)
            else:
                if node is not None and dst is None:
                    print '"%s"->"%s" [ label="%s", fontsize=10, color="%s" ];' % (src, node, label, color)
                elif node is None and dst is not None:
                    print '"%s"->"%s" [ label="%s", fontsize=10, color="%s" ];' % (src, dst, label, color)
                elif node is not None and dst is not None:
                    print '"%s"->"%s" [ label="%s", fontsize=10, color="%s" ];' % (src, node, "", color)
                    print '"%s"->"%s" [ label="%s", fontsize=10, color="%s" ];' % (node, dst, label, color)

        print '}'

if __name__ == "__main__":
    DotBot.from_command_line().execute()
<|endoftext|>


<commit_before>##
# Copyright (c) 2011 Apple 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
#
#     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 sqlparse
import os
from gzip import GzipFile
import collections
import tables
import textwrap
import sys
import getopt

def safePercent(x, y, multiplier=100):
    return ((multiplier * x) / y) if y else 0

def _is_literal(token):
    if token.ttype in sqlparse.tokens.Literal:
        return True
    if token.ttype == sqlparse.tokens.Keyword and token.value in (u'True', u'False'):
        return True
    return False

def _substitute(expression, replacement):
    try:
        expression.tokens
    except AttributeError:
        return

    for i, token in enumerate(expression.tokens):
        if _is_literal(token):
            expression.tokens[i] = replacement
        elif token.is_whitespace():
            expression.tokens[i] = sqlparse.sql.Token('Whitespace', ' ')
        else:
            _substitute(token, replacement)

def sqlnormalize(sql):
    try:
        statements = sqlparse.parse(sql)
    except ValueError, e:
        print e
    # Replace any literal values with placeholders
    qmark = sqlparse.sql.Token('Operator', '?')
    _substitute(statements[0], qmark)
    return sqlparse.format(statements[0].to_unicode().encode('ascii'))

COLUMN_userid = 0
COLUMN_dbid = 1
COLUMN_query = 2
COLUMN_calls = 3
COLUMN_total_time = 4
COLUMN_rows = 5
COLUMN_shared_blks_hit = 6
COLUMN_shared_blks_read = 7
COLUMN_shared_blks_written = 8
COLUMN_local_blks_hit = 9
COLUMN_local_blks_read = 10
COLUMN_local_blks_written = 11
COLUMN_temp_blks_read = 12
COLUMN_temp_blks_written = 13

def sqlStatementsReport(entries):

    dcount = collections.defaultdict(int)
    dtime = collections.defaultdict(float)
    for entry in entries:
        dcount[entry[COLUMN_query]] += int(entry[COLUMN_calls])
        dtime[entry[COLUMN_query]] += float(entry[COLUMN_total_time])

    daverage = {}
    for k in dcount.keys():
        daverage[k] = dtime[k]/dcount[k]

    counttotal = sum(dcount.values())
    timetotal = sum(dtime.values())
    averagetotal = sum(daverage.values())

    for sorttype, sortedkeys in (
        ("count", [i[0] for i in sorted(dcount.iteritems(), key=lambda x:x[1], reverse=True)],),
        ("total time", [i[0] for i in sorted(dtime.iteritems(), key=lambda x:x[1], reverse=True)],),
        ("average time", [i[0] for i in sorted(daverage.iteritems(), key=lambda x:x[1], reverse=True)],),
    ):
        table = tables.Table()
        table.addHeader(("Statement", "Count", "Count %", "Total Time", "Total Time %", "Av. Time", "Av. Time %",))
        table.setDefaultColumnFormats((
            tables.Table.ColumnFormat("%s", tables.Table.ColumnFormat.LEFT_JUSTIFY),
            tables.Table.ColumnFormat("%d", tables.Table.ColumnFormat.RIGHT_JUSTIFY),
            tables.Table.ColumnFormat("%.2f%%", tables.Table.ColumnFormat.RIGHT_JUSTIFY),
            tables.Table.ColumnFormat("%f", tables.Table.ColumnFormat.RIGHT_JUSTIFY),
            tables.Table.ColumnFormat("%.2f%%", tables.Table.ColumnFormat.RIGHT_JUSTIFY),
            tables.Table.ColumnFormat("%f", tables.Table.ColumnFormat.RIGHT_JUSTIFY),
            tables.Table.ColumnFormat("%.2f%%", tables.Table.ColumnFormat.RIGHT_JUSTIFY),
        ))

        for key in sortedkeys:

            keylines = textwrap.wrap(key, 72, subsequent_indent="  ")
            table.addRow((
                keylines[0],
                dcount[key],
                safePercent(dcount[key], counttotal, 100.0),
                dtime[key],
                safePercent(dtime[key], timetotal, 100.0),
                daverage[key],
                safePercent(daverage[key], averagetotal, 100.0),
            ))

            for keyline in keylines[1:]:
                table.addRow((
                    keyline,
                    None,
                    None,
                    None,
                    None,
                    None,
                    None,
                ))

        print "Queries sorted by %s" % (sorttype,)
        table.printTable()
        print ""

def parseStats(logFilePath, donormlize=True, verbose=False):

    fpath = os.path.expanduser(logFilePath)
    if fpath.endswith(".gz"):
        f = GzipFile(fpath)
    else:
        f = open(fpath)

    # Punt past data
    for line in f:
        if line.startswith("---"):
            break

    entries = []
    for line in f:
        bits = line.split("|")
        if len(bits) > COLUMN_query:
            while bits[COLUMN_query].endswith("+"):
                line = f.next()
                newbits = line.split("|")
                bits[COLUMN_query] = bits[COLUMN_query][:-1] + newbits[COLUMN_query]

            if donormlize:
                bits[COLUMN_query] = sqlnormalize(bits[COLUMN_query].strip())

            pos = bits[COLUMN_query].find("BEGIN:VCALENDAR")
            if pos != -1:
                bits[COLUMN_query] = bits[COLUMN_query][:pos]

            if bits[COLUMN_query] not in (
                "BEGIN",
                "COMMIT",
                "ROLLBACK",
            ) and bits[COLUMN_query].find("pg_catalog") == -1:
                bits = [bit.strip() for bit in bits]
                entries.append(bits)
                if verbose and divmod(len(entries), 1000)[1] == 0:
                    print "%d entries" % (len(entries),)
                #if float(bits[COLUMN_total_time]) > 1:
                #    print bits[COLUMN_total_time], bits[COLUMN_query]

    if verbose:
        print "Read %d entries" % (len(entries,))

    sqlStatementsReport(entries)

def usage(error_msg=None):
    if error_msg:
        print error_msg

    print """Usage: pg_stats_analysis.py [options] FILE
Options:
    -h             Print this help and exit
    -v             Generate progress information
    --no-normalize Do not normalize SQL statements

Arguments:
    FILE      File name for pg_stat_statements output to analyze.

Description:
This utility will analyze the output of s pg_stat_statement table.
"""

    if error_msg:
        raise ValueError(error_msg)
    else:
        sys.exit(0)

if __name__ == '__main__':

    normalize = True
    verbose = False
    options, args = getopt.getopt(sys.argv[1:], "hv", ["no-normalize",])

    for option, value in options:
        if option == "-h":
            usage()
        elif option == "-v":
            verbose = True
        elif option == "--no-normalize":
            normalize = False
        else:
            usage("Unrecognized option: %s" % (option,))

    # Process arguments
    if len(args) != 1:
        usage("Must have a file argument")

    parseStats(args[0], normalize, verbose)
<commit_msg>Better handling of truncated statements.<commit_after>##
# Copyright (c) 2011 Apple 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
#
#     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 sqlparse
import os
from gzip import GzipFile
import collections
import tables
import textwrap
import sys
import getopt

def safePercent(x, y, multiplier=100):
    return ((multiplier * x) / y) if y else 0

def _is_literal(token):
    if token.ttype in sqlparse.tokens.Literal:
        return True
    if token.ttype == sqlparse.tokens.Keyword and token.value in (u'True', u'False'):
        return True
    return False

def _substitute(expression, replacement):
    try:
        expression.tokens
    except AttributeError:
        return

    for i, token in enumerate(expression.tokens):
        if _is_literal(token):
            expression.tokens[i] = replacement
        elif token.is_whitespace():
            expression.tokens[i] = sqlparse.sql.Token('Whitespace', ' ')
        else:
            _substitute(token, replacement)

def sqlnormalize(sql):
    try:
        statements = sqlparse.parse(sql)
    except ValueError, e:
        print e
    # Replace any literal values with placeholders
    qmark = sqlparse.sql.Token('Operator', '?')
    _substitute(statements[0], qmark)
    return sqlparse.format(statements[0].to_unicode().encode('ascii'))

COLUMN_userid = 0
COLUMN_dbid = 1
COLUMN_query = 2
COLUMN_calls = 3
COLUMN_total_time = 4
COLUMN_rows = 5
COLUMN_shared_blks_hit = 6
COLUMN_shared_blks_read = 7
COLUMN_shared_blks_written = 8
COLUMN_local_blks_hit = 9
COLUMN_local_blks_read = 10
COLUMN_local_blks_written = 11
COLUMN_temp_blks_read = 12
COLUMN_temp_blks_written = 13

def sqlStatementsReport(entries):

    dcount = collections.defaultdict(int)
    dtime = collections.defaultdict(float)
    for entry in entries:
        dcount[entry[COLUMN_query]] += int(entry[COLUMN_calls])
        dtime[entry[COLUMN_query]] += float(entry[COLUMN_total_time])

    daverage = {}
    for k in dcount.keys():
        daverage[k] = dtime[k]/dcount[k]

    counttotal = sum(dcount.values())
    timetotal = sum(dtime.values())
    averagetotal = sum(daverage.values())

    for sorttype, sortedkeys in (
        ("count", [i[0] for i in sorted(dcount.iteritems(), key=lambda x:x[1], reverse=True)],),
        ("total time", [i[0] for i in sorted(dtime.iteritems(), key=lambda x:x[1], reverse=True)],),
        ("average time", [i[0] for i in sorted(daverage.iteritems(), key=lambda x:x[1], reverse=True)],),
    ):
        table = tables.Table()
        table.addHeader(("Statement", "Count", "Count %", "Total Time", "Total Time %", "Av. Time", "Av. Time %",))
        table.setDefaultColumnFormats((
            tables.Table.ColumnFormat("%s", tables.Table.ColumnFormat.LEFT_JUSTIFY),
            tables.Table.ColumnFormat("%d", tables.Table.ColumnFormat.RIGHT_JUSTIFY),
            tables.Table.ColumnFormat("%.2f%%", tables.Table.ColumnFormat.RIGHT_JUSTIFY),
            tables.Table.ColumnFormat("%f", tables.Table.ColumnFormat.RIGHT_JUSTIFY),
            tables.Table.ColumnFormat("%.2f%%", tables.Table.ColumnFormat.RIGHT_JUSTIFY),
            tables.Table.ColumnFormat("%f", tables.Table.ColumnFormat.RIGHT_JUSTIFY),
            tables.Table.ColumnFormat("%.2f%%", tables.Table.ColumnFormat.RIGHT_JUSTIFY),
        ))

        for key in sortedkeys:

            keylines = textwrap.wrap(key, 72, subsequent_indent="  ")
            table.addRow((
                keylines[0],
                dcount[key],
                safePercent(dcount[key], counttotal, 100.0),
                dtime[key],
                safePercent(dtime[key], timetotal, 100.0),
                daverage[key],
                safePercent(daverage[key], averagetotal, 100.0),
            ))

            for keyline in keylines[1:]:
                table.addRow((
                    keyline,
                    None,
                    None,
                    None,
                    None,
                    None,
                    None,
                ))

        print "Queries sorted by %s" % (sorttype,)
        table.printTable()
        print ""

def parseStats(logFilePath, donormlize=True, verbose=False):

    fpath = os.path.expanduser(logFilePath)
    if fpath.endswith(".gz"):
        f = GzipFile(fpath)
    else:
        f = open(fpath)

    # Punt past data
    for line in f:
        if line.startswith("---"):
            break

    entries = []
    for line in f:
        bits = line.split("|")
        if len(bits) > COLUMN_query:
            while bits[COLUMN_query].endswith("+"):
                line = f.next()
                newbits = line.split("|")
                bits[COLUMN_query] = bits[COLUMN_query][:-1] + newbits[COLUMN_query]

            pos = bits[COLUMN_query].find("BEGIN:VCALENDAR")
            if pos != -1:
                bits[COLUMN_query] = bits[COLUMN_query][:pos]

            if donormlize:
                bits[COLUMN_query] = sqlnormalize(bits[COLUMN_query].strip())

            if bits[COLUMN_query] not in (
                "BEGIN",
                "COMMIT",
                "ROLLBACK",
            ) and bits[COLUMN_query].find("pg_catalog") == -1:
                bits = [bit.strip() for bit in bits]
                entries.append(bits)
                if verbose and divmod(len(entries), 1000)[1] == 0:
                    print "%d entries" % (len(entries),)
                #if float(bits[COLUMN_total_time]) > 1:
                #    print bits[COLUMN_total_time], bits[COLUMN_query]

    if verbose:
        print "Read %d entries" % (len(entries,))

    sqlStatementsReport(entries)

def usage(error_msg=None):
    if error_msg:
        print error_msg

    print """Usage: pg_stats_analysis.py [options] FILE
Options:
    -h             Print this help and exit
    -v             Generate progress information
    --no-normalize Do not normalize SQL statements

Arguments:
    FILE      File name for pg_stat_statements output to analyze.

Description:
This utility will analyze the output of s pg_stat_statement table.
"""

    if error_msg:
        raise ValueError(error_msg)
    else:
        sys.exit(0)

if __name__ == '__main__':

    normalize = True
    verbose = False
    options, args = getopt.getopt(sys.argv[1:], "hv", ["no-normalize",])

    for option, value in options:
        if option == "-h":
            usage()
        elif option == "-v":
            verbose = True
        elif option == "--no-normalize":
            normalize = False
        else:
            usage("Unrecognized option: %s" % (option,))

    # Process arguments
    if len(args) != 1:
        usage("Must have a file argument")

    parseStats(args[0], normalize, verbose)
<|endoftext|>


<commit_before>"""
Example module
"""
import logging

from telegram import Bot, Update
from telegram.ext import MessageHandler, Updater, Filters

import constants # pylint: disable=E0401
LOGGER = logging.getLogger("Example Module")
def void(*_):
    """
    Simple function that can accept anything...
    ...and do nothing
    """
    return

def preload(updater: Updater, level):
    """
    This loads whenever plugin starts
    Even if you dont need it, you SHOULD put at least
    return None, otherwise your plugin wont load
    """
    LOGGER.debug("TEST")
    # Here is example of adding your own MessageHandler
    # You SHOULD specify group=level, otherwise bot may not work correctly
    updater.dispatcher.add_handler(MessageHandler(Filters.text, void), group=level)

def helloworld(bot: Bot, update: Update, user, args): # pylint: disable=W0613
    """/helloworld command"""
    return "Hello World", constants.TEXT

COMMANDS = [
    {
        "command":"/helloworld",
        "function":helloworld,
        "description":"Sends Hello World!",
        "inline_support":True
    }
]
<commit_msg>Removed old example plugin<commit_after><|endoftext|>


<commit_before>import numpy as np
import tensorflow as tf

from tensor_train import TensorTrain
import ops
import initializers
import riemannian

class RiemannianTest(tf.test.TestCase):

  def testProjectOnItself(self):
    # Projection of X into the tangent space of itself is X: P_x(x) = x.
    tens = initializers.random_tensor((2, 3, 4))
    proj = riemannian.project(tens, tens)
    with self.test_session() as sess:
      actual_val, desired_val = sess.run((ops.full(proj), ops.full(tens)))
      self.assertAllClose(desired_val, actual_val)

  def testProject(self):
    # Compare our projection with the results obtained (and precomputed) from
    # tt.riemannian.project which is well tested.
    tangent_tens_cores = ([[[-0.42095269,  0.02130842],
         [-0.4181081 ,  0.42945687],
         [ 0.45972439, -0.4525616 ],
         [-0.17159869, -0.14505528]]], [[[ 0.23344421],
         [ 0.81480049],
         [-0.92385135]],

        [[-0.19279465],
         [ 0.524976  ],
         [-0.40149197]]])
    tangent_tens = TensorTrain(tangent_tens_cores, (4, 3), (1, 2, 1))
    tens_cores = ([[[-1.01761142,  0.36075896, -0.2493624 ],
         [-0.99896565, -1.12685474,  1.02832458],
         [ 1.08739724, -0.6537435 ,  1.99975537],
         [ 0.35128005,  0.40395104, -0.16790072]]], [[[ 0.34105142],
         [ 0.10371947],
         [-1.76464904]],

        [[ 0.32639758],
         [-1.27843174],
         [-1.41590327]],

        [[ 0.76616274],
         [ 0.6577514 ],
         [ 2.13703185]]])
    tens = TensorTrain(tens_cores, (4, 3), (1, 3, 1))
    desired_projection = [[-0.67638254, -1.17163914,  0.29850939],
       [-1.66479093, -0.99003251,  2.46629195],
       [-0.04847773, -0.72908174,  0.20142675],
       [ 0.34431125, -0.20935516, -1.15864246]]
    proj = riemannian.project(tangent_tens, tens)
    with self.test_session() as sess:
      self.assertAllClose(desired_projection, ops.full(proj).eval())

  def testProjectSum(self):
    # Test projecting a batch of TT-tensors.
    tens = initializers.random_tensor_batch((2, 3, 4), batch_size=3)
    tangent_tens = initializers.random_tensor((2, 3, 4), 3)
    weighted_sum = tens[0] + tens[1] + tens[2]
    direct_proj = riemannian.project(tangent_tens, weighted_sum)
    actual_proj = riemannian.project(tangent_tens, tens)
    with self.test_session() as sess:
      res = sess.run((ops.full(direct_proj), ops.full(actual_proj)))
      desired_val, actual_val = res
      self.assertAllClose(desired_val, actual_val)

  def testProjectSumCoef(self):
    # Test projecting a batch of TT-tensors with providing coefs.
    tens = initializers.random_tensor_batch((2, 3, 4), 3, batch_size=4)
    coef = [0.1, -2, 0, 0.4]
    tangent_tens = initializers.random_tensor((2, 3, 4), 4)
    weighted_sum = coef[0] * tens[0] + coef[1] * tens[1] + coef[2] * tens[2]
    weighted_sum += coef[3] * tens[3]
    direct_proj = riemannian.project(tangent_tens, weighted_sum)
    actual_proj = riemannian.project(tangent_tens, tens, coef)
    with self.test_session() as sess:
      res = sess.run((ops.full(direct_proj), ops.full(actual_proj)))
      desired_val, actual_val = res
      self.assertAllClose(desired_val, actual_val)

  def testProjectMatrixOnItself(self):
    # Project a TT-matrix on itself.
    # Projection of X into the tangent space of itself is X: P_x(x) = x.
    tt_mat = initializers.random_matrix(((2, 3, 4), (2, 3, 4)))
    proj = riemannian.project(tt_mat, tt_mat)
    with self.test_session() as sess:
      actual_val, desired_val = sess.run((ops.full(proj), ops.full(tt_mat)))
      self.assertAllClose(desired_val, actual_val)

if __name__ == "__main__":
  tf.test.main()
<commit_msg>Reorder project arguments for more intuitive API<commit_after>import numpy as np
import tensorflow as tf

from tensor_train import TensorTrain
import ops
import initializers
import riemannian

class RiemannianTest(tf.test.TestCase):

  def testProjectOnItself(self):
    # Projection of X into the tangent space of itself is X: P_x(x) = x.
    tens = initializers.random_tensor((2, 3, 4))
    proj = riemannian.project(tens, tens)
    with self.test_session() as sess:
      actual_val, desired_val = sess.run((ops.full(proj), ops.full(tens)))
      self.assertAllClose(desired_val, actual_val)

  def testProject(self):
    # Compare our projection with the results obtained (and precomputed) from
    # tt.riemannian.project which is well tested.
    tangent_tens_cores = ([[[-0.42095269,  0.02130842],
         [-0.4181081 ,  0.42945687],
         [ 0.45972439, -0.4525616 ],
         [-0.17159869, -0.14505528]]], [[[ 0.23344421],
         [ 0.81480049],
         [-0.92385135]],

        [[-0.19279465],
         [ 0.524976  ],
         [-0.40149197]]])
    tangent_tens = TensorTrain(tangent_tens_cores, (4, 3), (1, 2, 1))
    tens_cores = ([[[-1.01761142,  0.36075896, -0.2493624 ],
         [-0.99896565, -1.12685474,  1.02832458],
         [ 1.08739724, -0.6537435 ,  1.99975537],
         [ 0.35128005,  0.40395104, -0.16790072]]], [[[ 0.34105142],
         [ 0.10371947],
         [-1.76464904]],

        [[ 0.32639758],
         [-1.27843174],
         [-1.41590327]],

        [[ 0.76616274],
         [ 0.6577514 ],
         [ 2.13703185]]])
    tens = TensorTrain(tens_cores, (4, 3), (1, 3, 1))
    desired_projection = [[-0.67638254, -1.17163914,  0.29850939],
       [-1.66479093, -0.99003251,  2.46629195],
       [-0.04847773, -0.72908174,  0.20142675],
       [ 0.34431125, -0.20935516, -1.15864246]]
    proj = riemannian.project(tens, tangent_tens)
    with self.test_session() as sess:
      self.assertAllClose(desired_projection, ops.full(proj).eval())

  def testProjectSum(self):
    # Test projecting a batch of TT-tensors.
    tens = initializers.random_tensor_batch((2, 3, 4), batch_size=3)
    tangent_tens = initializers.random_tensor((2, 3, 4), 3)
    weighted_sum = tens[0] + tens[1] + tens[2]
    direct_proj = riemannian.project(weighted_sum, tangent_tens)
    actual_proj = riemannian.project(tens, tangent_tens)
    with self.test_session() as sess:
      res = sess.run((ops.full(direct_proj), ops.full(actual_proj)))
      desired_val, actual_val = res
      self.assertAllClose(desired_val, actual_val)

  def testProjectSumCoef(self):
    # Test projecting a batch of TT-tensors with providing coefs.
    tens = initializers.random_tensor_batch((2, 3, 4), 3, batch_size=4)
    coef = [0.1, -2, 0, 0.4]
    tangent_tens = initializers.random_tensor((2, 3, 4), 4)
    weighted_sum = coef[0] * tens[0] + coef[1] * tens[1] + coef[2] * tens[2]
    weighted_sum += coef[3] * tens[3]
    direct_proj = riemannian.project(weighted_sum, tangent_tens)
    actual_proj = riemannian.project(tens, tangent_tens, coef)
    with self.test_session() as sess:
      res = sess.run((ops.full(direct_proj), ops.full(actual_proj)))
      desired_val, actual_val = res
      self.assertAllClose(desired_val, actual_val)

  def testProjectMatrixOnItself(self):
    # Project a TT-matrix on itself.
    # Projection of X into the tangent space of itself is X: P_x(x) = x.
    tt_mat = initializers.random_matrix(((2, 3, 4), (2, 3, 4)))
    proj = riemannian.project(tt_mat, tt_mat)
    with self.test_session() as sess:
      actual_val, desired_val = sess.run((ops.full(proj), ops.full(tt_mat)))
      self.assertAllClose(desired_val, actual_val)

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


<commit_before>"""."""

import numpy as np
from pymodels import tb, ts
import pyaccel

kick = 1e-5

def get_tb_posang_respm(model=None):
    """Return TB position and angle correction matrix."""
    if model is None:
        model, _ = tb.create_accelerator()
    fam_data = tb.families.get_family_data(model)

    ch = [fam_data['CH']['index'][-1], fam_data['InjSept']['index'][0]]
    cv = [fam_data['CV']['index'][-2], fam_data['CV']['index'][-1]]

    return _calc_posang_matrices(model, ch, cv)

def get_ts_posang_respm(corrs_type='CH-Sept', model=None):
    """Return TS position and angle correction matrix."""
    if model is None:
        model, _ = ts.create_accelerator()
    fam_data = ts.families.get_family_data(model)

    if corrs_type == 'CH-Sept':
        ch = [fam_data['CH']['index'][-1],
              fam_data['InjSeptF']['index'][0]]
    else:
        ch = [fam_data['InjSeptG']['index'][0],
              fam_data['InjSeptG']['index'][1],
              fam_data['InjSeptF']['index'][0]]
    cv = [fam_data['CV']['index'][-2], fam_data['CV']['index'][-1]]

    return _calc_posang_matrices(model, ch, cv)

def _calc_posang_matrices(model, ch, cv):
    ini = ch[0][0]
    mat_h_aux = np.zeros((2, len(ch)))
    for i in range(len(ch)):
        pyaccel.lattice.set_attribute(
            model, 'hkick_polynom', ch[i], kick/2)
        p1 = pyaccel.tracking.line_pass(model[ini:len(model)], 6*[0])
        pyaccel.lattice.set_attribute(
            model, 'hkick_polynom', ch[i], -kick/2)
        p2 = pyaccel.tracking.line_pass(model[ini:len(model)], 6*[0])
        pyaccel.lattice.set_attribute(
            model, 'hkick_polynom', ch[i], 0)
        mat_h_aux[:, i] = (p1[0][0:2] - p2[0][0:2])/kick

    if len(ch) == 3:
        mat_h = [[mat_h_aux[0, 0]+mat_h_aux[0, 1], mat_h_aux[0, 2]],
                 [mat_h_aux[1, 0]+mat_h_aux[1, 1], mat_h_aux[1, 2]]]
    else:
        mat_h = mat_h_aux

    ini = cv[0][0]
    mat_v = np.zeros((2, 2))
    for i in range(len(cv)):
        pyaccel.lattice.set_attribute(
            model, 'vkick_polynom', cv[i], kick/2)
        p1 = pyaccel.tracking.line_pass(model[ini:len(model)], 6*[0])
        pyaccel.lattice.set_attribute(
            model, 'vkick_polynom', cv[i], -kick/2)
        p2 = pyaccel.tracking.line_pass(model[ini:len(model)], 6*[0])
        pyaccel.lattice.set_attribute(
            model, 'vkick_polynom', cv[i], 0)
        mat_v[:, i] = (p1[0][2:4] - p2[0][2:4])/kick

    return mat_h, mat_v
<commit_msg>COMMISPOSANG.BUG: divide kickangle by the number of segments of the corrector.<commit_after>"""."""

import numpy as np
from pymodels import tb, ts
import pyaccel

kick = 1e-5

def get_tb_posang_respm(model=None):
    """Return TB position and angle correction matrix."""
    if model is None:
        model, _ = tb.create_accelerator()
    fam_data = tb.families.get_family_data(model)

    ch = [fam_data['CH']['index'][-1], fam_data['InjSept']['index'][0]]
    cv = [fam_data['CV']['index'][-2], fam_data['CV']['index'][-1]]

    return _calc_posang_matrices(model, ch, cv)

def get_ts_posang_respm(corrs_type='CH-Sept', model=None):
    """Return TS position and angle correction matrix."""
    if model is None:
        model, _ = ts.create_accelerator()
    fam_data = ts.families.get_family_data(model)

    if corrs_type == 'CH-Sept':
        ch = [fam_data['CH']['index'][-1],
              fam_data['InjSeptF']['index'][0]]
    else:
        ch = [fam_data['InjSeptG']['index'][0],
              fam_data['InjSeptG']['index'][1],
              fam_data['InjSeptF']['index'][0]]
    cv = [fam_data['CV']['index'][-2], fam_data['CV']['index'][-1]]

    return _calc_posang_matrices(model, ch, cv)

def _calc_posang_matrices(model, ch, cv):
    ini = ch[0][0]
    mat_h_aux = np.zeros((2, len(ch)))
    for i in range(len(ch)):
        pyaccel.lattice.set_attribute(
            model, 'hkick_polynom', ch[i], kick/2/len(ch[i]))
        p1 = pyaccel.tracking.line_pass(model[ini:len(model)], 6*[0])
        pyaccel.lattice.set_attribute(
            model, 'hkick_polynom', ch[i], -kick/2/len(ch[i]))
        p2 = pyaccel.tracking.line_pass(model[ini:len(model)], 6*[0])
        pyaccel.lattice.set_attribute(
            model, 'hkick_polynom', ch[i], 0)
        mat_h_aux[:, i] = (p1[0][0:2] - p2[0][0:2])/kick

    if len(ch) == 3:
        mat_h = [[mat_h_aux[0, 0]+mat_h_aux[0, 1], mat_h_aux[0, 2]],
                 [mat_h_aux[1, 0]+mat_h_aux[1, 1], mat_h_aux[1, 2]]]
    else:
        mat_h = mat_h_aux

    ini = cv[0][0]
    mat_v = np.zeros((2, 2))
    for i in range(len(cv)):
        pyaccel.lattice.set_attribute(
            model, 'vkick_polynom', cv[i], kick/2/len(cv[i]))
        p1 = pyaccel.tracking.line_pass(model[ini:len(model)], 6*[0])
        pyaccel.lattice.set_attribute(
            model, 'vkick_polynom', cv[i], -kick/2/len(cv[i]))
        p2 = pyaccel.tracking.line_pass(model[ini:len(model)], 6*[0])
        pyaccel.lattice.set_attribute(
            model, 'vkick_polynom', cv[i], 0)
        mat_v[:, i] = (p1[0][2:4] - p2[0][2:4])/kick

    return mat_h, mat_v
<|endoftext|>


<commit_before>import datetime
import os
import re
import socket
import struct
import subprocess
import time

try:
    import pcapy
except ImportError:
    if subprocess.mswindows:
        exit("[!] please install Pcapy (e.g. 'https://breakingcode.wordpress.com/?s=pcapy') and WinPcap (e.g. 'http://www.winpcap.org/install/')")
    else:
        exit("[!] please install Pcapy (e.g. 'apt-get install python-pcapy')")

_auxiliary = {}
_cap = None
_datalink = None
_traffic = {}

CAPTURE_FILTER = "(tcp[13] == 2) or not tcp"  # SYN or not TCP
MONITOR_INTERFACE = "any"
SNAP_LEN = 100
IPPROTO = 8
ETH_LENGTH = 14
LAST_FILENAME = None
LAST_WRITE = None
WRITE_PERIOD = 60 * 15
IPPROTO_LUT = dict(((getattr(socket, _), _.replace("IPPROTO_", "")) for _ in dir(socket) if _.startswith("IPPROTO_")))
LOCAL_ADDRESSES = []
BLACKLISTED_ADDRESSES = ("255.255.255.255", "127.0.0.1", "0.0.0.0")
DATE_FORMAT = "%Y-%m-%d"
RESULTS_DIRECTORY = os.path.normpath(os.path.join(os.path.dirname(__file__), "./results"))

def _log_write(force=False, filename=None):
    global LAST_FILENAME
    global LAST_WRITE

    current = time.time()
    filename = filename or os.path.join(RESULTS_DIRECTORY, "%s.csv" % datetime.date.today().strftime(DATE_FORMAT))

    if LAST_WRITE is None:
        LAST_WRITE = current

    if LAST_FILENAME is None:
        LAST_FILENAME = filename

    if force or (current - LAST_WRITE) > WRITE_PERIOD:
        if not os.path.isdir(RESULTS_DIRECTORY):
            os.makedirs(RESULTS_DIRECTORY)

        with open(filename, "w+b") as f:
            results = []
            f.write("proto dst_port src_ip dst_ip timestamp\n")

            for key in _traffic:
                proto, dst_ip, dst_port = key.split(":")
                for src_ip in sorted(_traffic[key]):
                    sec = _auxiliary["%s:%s" % (key, src_ip)]
                    results.append((sec, (proto, dst_port, src_ip, dst_ip)))

            for sec, entry in sorted(results):
                f.write("%s %s\n" % (" ".join(str(_) for _ in entry), sec))

        LAST_WRITE = current

    if LAST_FILENAME != filename:
        if not force and LAST_WRITE != current:
            _log_write(True, LAST_FILENAME)

        LAST_FILENAME = filename

        _traffic.clear()
        _auxiliary.clear()

def _process_packet(packet, sec, usec):
    try:
        if _datalink == pcapy.DLT_LINUX_SLL:
            packet = packet[2:]

        eth_header = struct.unpack("!HH8sH", packet[:ETH_LENGTH])
        eth_protocol = socket.ntohs(eth_header[3])

        if eth_protocol == IPPROTO:  # IP
            ip_header = struct.unpack("!BBHHHBBH4s4s", packet[ETH_LENGTH:ETH_LENGTH + 20])
            ip_length = ip_header[2]
            packet = packet[:ETH_LENGTH + ip_length]  # truncate
            iph_length = (ip_header[0] & 0xF) << 2
            protocol = ip_header[6]
            src_ip = socket.inet_ntoa(ip_header[8])
            dst_ip = socket.inet_ntoa(ip_header[9])

            proto = IPPROTO_LUT.get(protocol)

            if proto is None:
                return

            # only process SYN packets
            if protocol == socket.IPPROTO_TCP:  # TCP
                i = iph_length + ETH_LENGTH
                src_port, dst_port, _, _, doff_reserved, flags = struct.unpack("!HHLLBB", packet[i:i + 14])
                if flags == 2:  # SYN set (only)
                    local_src = False
                    for prefix, mask in LOCAL_ADDRESSES:
                        if addr_to_int(src_ip) & mask == prefix:
                            local_src = True
                            break
                    if not local_src and dst_ip not in BLACKLISTED_ADDRESSES:
                        key = "%s:%s:%s" % (proto, dst_ip, dst_port)
                        if key not in _traffic:
                            _traffic[key] = set()
                        _traffic[key].add(src_ip)
                        _auxiliary["%s:%s" % (key, src_ip)] = sec

            else:
                if protocol == socket.IPPROTO_UDP:  # UDP
                    i = iph_length + ETH_LENGTH
                    _ = packet[i:i + 4]
                    if len(_) < 4:
                        return

                    src_port, dst_port = struct.unpack("!HH", _)
                else:                               # non-TCP/UDP (e.g. ICMP)
                    src_port, dst_port = '-', '-'

                flow = tuple(sorted((addr_to_int(src_ip), src_port, addr_to_int(dst_ip), dst_port)))
                if flow not in _auxiliary:
                    _auxiliary[flow] = True
                    local_src = False
                    for prefix, mask in LOCAL_ADDRESSES:
                        if addr_to_int(src_ip) & mask == prefix:
                            local_src = True
                            break
                    if not local_src and dst_ip not in BLACKLISTED_ADDRESSES:
                        key = "%s:%s:%s" % (proto, dst_ip, dst_port)
                        if key not in _traffic:
                            _traffic[key] = set()
                        _traffic[key].add(src_ip)
                        _auxiliary["%s:%s" % (key, src_ip)] = sec

    except KeyboardInterrupt:
        raise

    except Exception:
        pass

    finally:
        _log_write()

def packet_handler(header, packet):
    try:
        sec, usec = header.getts()
        _process_packet(packet, sec, usec)
    except socket.timeout:
        pass

def addr_to_int(value):
    _ = value.split('.')
    return (long(_[0]) << 24) + (long(_[1]) << 16) + (long(_[2]) << 8) + long(_[3])

def int_to_addr(value):
    return '.'.join(str(value >> n & 0xFF) for n in (24, 16, 8, 0))

def make_mask(bits):
    return 0xffffffff ^ (1 << 32 - bits) - 1

def local_addresses():
    items = []

    for cmd, regex in (("ifconfig", r"inet addr:([\d.]+) .*Mask:([\d.]+)"), ("ipconfig", r"IPv4 Address[^\n]+([\d.]+)\s+Subnet Mask[^\n]+([\d.]+)")):
        try:
            items = re.findall(regex, subprocess.check_output(cmd))
            break
        except:
            pass

    for ip, mask in items:
        LOCAL_ADDRESSES.append((addr_to_int(ip) & addr_to_int(mask), addr_to_int(mask)))

def main():
    global _cap
    global _datalink

    local_addresses()

    print "[i] opening interface '%s'" % MONITOR_INTERFACE
    try:
        _cap = pcapy.open_live(MONITOR_INTERFACE, SNAP_LEN, True, 0)
    except socket.error, ex:
        if "permitted" in str(ex):
            exit("\n[x] please run with sudo/Administrator privileges")
        elif "No such device" in str(ex):
            exit("\n[x] no such device '%s'" % MONITOR_INTERFACE)
        else:
            raise

    if CAPTURE_FILTER:
        print "[i] setting filter '%s'" % CAPTURE_FILTER
        _cap.setfilter(CAPTURE_FILTER)

    _datalink = _cap.datalink()
    if _datalink not in (pcapy.DLT_EN10MB, pcapy.DLT_LINUX_SLL):
        exit("[x] datalink type '%s' not supported" % _datalink)

    try:
        _cap.loop(-1, packet_handler)
    except KeyboardInterrupt:
        print "\r[x] Ctrl-C pressed"
        _log_write(True)

if __name__ == "__main__":
    main()
<commit_msg>Minor improvements<commit_after>import datetime
import os
import re
import socket
import struct
import subprocess
import time

try:
    import pcapy
except ImportError:
    if subprocess.mswindows:
        exit("[!] please install Pcapy (e.g. 'https://breakingcode.wordpress.com/?s=pcapy') and WinPcap (e.g. 'http://www.winpcap.org/install/')")
    else:
        exit("[!] please install Pcapy (e.g. 'apt-get install python-pcapy')")

_auxiliary = {}
_cap = None
_datalink = None
_traffic = {}

CAPTURE_FILTER = "(tcp[13] == 2) or not tcp"  # SYN or not TCP
MONITOR_INTERFACE = "any"
SNAP_LEN = 100
IPPROTO = 8
ETH_LENGTH = 14
LAST_FILENAME = None
LAST_WRITE = None
WRITE_PERIOD = 60 * 15
IPPROTO_LUT = dict(((getattr(socket, _), _.replace("IPPROTO_", "")) for _ in dir(socket) if _.startswith("IPPROTO_")))
LOCAL_ADDRESSES = []
BLACKLISTED_ADDRESSES = ("255.255.255.255", "127.0.0.1", "0.0.0.0")
DATE_FORMAT = "%Y-%m-%d"
RESULTS_DIRECTORY = os.path.normpath(os.path.join(os.path.dirname(__file__), "./results"))
CSV_HEADER = "proto dst_port dst_ip src_ip first_seen last_seen count"

def _log_write(force=False, filename=None):
    global LAST_FILENAME
    global LAST_WRITE

    current = time.time()
    filename = filename or os.path.join(RESULTS_DIRECTORY, "%s.csv" % datetime.date.today().strftime(DATE_FORMAT))

    if LAST_WRITE is None:
        LAST_WRITE = current

    if LAST_FILENAME is None:
        LAST_FILENAME = filename

    if force or (current - LAST_WRITE) > WRITE_PERIOD:
        if not os.path.isdir(RESULTS_DIRECTORY):
            os.makedirs(RESULTS_DIRECTORY)

        with open(filename, "w+b") as f:
            results = []
            f.write("%s\n" % CSV_HEADER)

            for dst_key in _traffic:
                proto, dst_ip, dst_port = dst_key.split(":")
                for src_ip in _traffic[dst_key]:
                    stat_key = "%s:%s" % (dst_key, src_ip)
                    first_seen, last_seen, count = _auxiliary[stat_key]
                    results.append((proto, int(dst_port), dst_ip, src_ip, first_seen, last_seen, count))

            for entry in sorted(results):
                f.write("%s\n" % " ".join(str(_) for _ in entry))

        LAST_WRITE = current

    if LAST_FILENAME != filename:
        if not force and LAST_WRITE != current:
            _log_write(True, LAST_FILENAME)

        LAST_FILENAME = filename

        _traffic.clear()
        _auxiliary.clear()

def _process_packet(packet, sec, usec):
    try:
        if _datalink == pcapy.DLT_LINUX_SLL:
            packet = packet[2:]

        eth_header = struct.unpack("!HH8sH", packet[:ETH_LENGTH])
        eth_protocol = socket.ntohs(eth_header[3])

        if eth_protocol == IPPROTO:  # IP
            ip_header = struct.unpack("!BBHHHBBH4s4s", packet[ETH_LENGTH:ETH_LENGTH + 20])
            ip_length = ip_header[2]
            packet = packet[:ETH_LENGTH + ip_length]  # truncate
            iph_length = (ip_header[0] & 0xF) << 2

            protocol = ip_header[6]
            src_ip = socket.inet_ntoa(ip_header[8])
            dst_ip = socket.inet_ntoa(ip_header[9])

            proto = IPPROTO_LUT.get(protocol)

            local_src = False
            for prefix, mask in LOCAL_ADDRESSES:
                if addr_to_int(src_ip) & mask == prefix:
                    local_src = True
                    break

            if proto is None or local_src or any(_ in BLACKLISTED_ADDRESSES for _ in (src_ip, dst_ip)):
                return

            # only process SYN packets
            if protocol == socket.IPPROTO_TCP:      # TCP
                i = iph_length + ETH_LENGTH
                src_port, dst_port, _, _, _, flags = struct.unpack("!HHLLBB", packet[i:i + 14])

                dst_key = "%s:%s:%s" % (proto, dst_ip, dst_port)
                stat_key = "%s:%s" % (dst_key, src_ip)

                if flags == 2:                      # SYN set (only)
                    if dst_key not in _traffic:
                        _traffic[dst_key] = set()

                    _traffic[dst_key].add(src_ip)

                    if stat_key not in _auxiliary:
                        _auxiliary[stat_key] = [sec, sec, 1]
                    else:
                        _auxiliary[stat_key][1] = sec
                        _auxiliary[stat_key][2] += 1

            else:
                if protocol == socket.IPPROTO_UDP:  # UDP
                    i = iph_length + ETH_LENGTH
                    _ = packet[i:i + 4]
                    if len(_) < 4:
                        return

                    src_port, dst_port = struct.unpack("!HH", _)
                else:                               # non-TCP/UDP (e.g. ICMP)
                    src_port, dst_port = '-', '-'

                dst_key = "%s:%s:%s" % (proto, dst_ip, dst_port)
                stat_key = "%s:%s" % (dst_key, src_ip)

                flow = tuple(sorted((addr_to_int(src_ip), src_port, addr_to_int(dst_ip), dst_port)))

                if flow not in _auxiliary:
                    _auxiliary[flow] = True

                    if dst_key not in _traffic:
                        _traffic[dst_key] = set()

                    _traffic[dst_key].add(src_ip)
                    _auxiliary[stat_key] = [sec, sec, 1]

                elif stat_key in _auxiliary:
                    _auxiliary[stat_key][1] = sec
                    _auxiliary[stat_key][2] += 1

    except KeyboardInterrupt:
        raise

    except Exception:
        pass

    finally:
        _log_write()

def packet_handler(header, packet):
    try:
        sec, usec = header.getts()
        _process_packet(packet, sec, usec)
    except socket.timeout:
        pass

def addr_to_int(value):
    _ = value.split('.')
    return (long(_[0]) << 24) + (long(_[1]) << 16) + (long(_[2]) << 8) + long(_[3])

def int_to_addr(value):
    return '.'.join(str(value >> n & 0xFF) for n in (24, 16, 8, 0))

def make_mask(bits):
    return 0xffffffff ^ (1 << 32 - bits) - 1

def local_addresses():
    items = []

    for cmd, regex in (("ifconfig", r"inet addr:([\d.]+) .*Mask:([\d.]+)"), ("ipconfig", r"IPv4 Address[^\n]+([\d.]+)\s+Subnet Mask[^\n]+([\d.]+)")):
        try:
            items = re.findall(regex, subprocess.check_output(cmd))
            break
        except:
            pass

    for ip, mask in items:
        LOCAL_ADDRESSES.append((addr_to_int(ip) & addr_to_int(mask), addr_to_int(mask)))

def main():
    global _cap
    global _datalink

    local_addresses()

    print "[i] opening interface '%s'" % MONITOR_INTERFACE
    try:
        _cap = pcapy.open_live(MONITOR_INTERFACE, SNAP_LEN, True, 0)
    except socket.error, ex:
        if "permitted" in str(ex):
            exit("\n[x] please run with sudo/Administrator privileges")
        elif "No such device" in str(ex):
            exit("\n[x] no such device '%s'" % MONITOR_INTERFACE)
        else:
            raise
    except Exception, ex:
        if "Operation not permitted" in str(ex):
            exit("[x] please run with sudo/Administrator privileges")
        else:
            raise

    if CAPTURE_FILTER:
        print "[i] setting filter '%s'" % CAPTURE_FILTER
        _cap.setfilter(CAPTURE_FILTER)

    _datalink = _cap.datalink()
    if _datalink not in (pcapy.DLT_EN10MB, pcapy.DLT_LINUX_SLL):
        exit("[x] datalink type '%s' not supported" % _datalink)

    try:
        _cap.loop(-1, packet_handler)
    except:
        print "\r[x] stopped"
    finally:
        _log_write(True)

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