<commit_before>"""
    gui/filter
    ~~~~~~~~~~~~~~~~~~~~

    Graphical user interface for filtering localization lists

    :author: Joerg Schnitzbauer, 2015
"""

import sys
import traceback
from PyQt4 import QtCore, QtGui
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg, NavigationToolbar2QT
import matplotlib.pyplot as plt
from matplotlib.widgets import SpanSelector, RectangleSelector
from matplotlib.colors import LogNorm
import numpy as np
import os.path
from picasso import io, lib

plt.style.use('ggplot')

class TableModel(QtCore.QAbstractTableModel):

    def __init__(self, locs, parent=None):
        super().__init__(parent)
        self.locs = locs

    def columnCount(self, parent):
        try:
            return len(self.locs[0])
        except IndexError:
            return 0

    def rowCount(self, parent):
        return self.locs.shape[0]

    def data(self, index, role=QtCore.Qt.DisplayRole):
        if role == QtCore.Qt.DisplayRole:
            data = self.locs[index.row()][index.column()]
            return str(data)
        return None

    def headerData(self, section, orientation, role):
        if role == QtCore.Qt.DisplayRole:
            if orientation == QtCore.Qt.Horizontal:
                return self.locs.dtype.names[section]
            elif orientation == QtCore.Qt.Vertical:
                return section
        return None

class TableView(QtGui.QTableView):

    def __init__(self, window, parent=None):
        super().__init__(parent)
        self.window = window

    def dragEnterEvent(self, event):
        if event.mimeData().hasUrls():
            event.accept()
        else:
            event.ignore()

    def dragMoveEvent(self, event):
        event.accept()

    def dropEvent(self, event):
        urls = event.mimeData().urls()
        path = urls[0].toLocalFile()
        extension = os.path.splitext(path)[1].lower()
        if extension == '.hdf5':
            self.window.open(path)

class PlotWindow(QtGui.QWidget):

    def __init__(self, main_window, locs):
        super().__init__()
        self.main_window = main_window
        self.locs = locs
        self.figure = plt.Figure()
        self.canvas = FigureCanvasQTAgg(self.figure)
        self.plot()
        vbox = QtGui.QVBoxLayout()
        self.setLayout(vbox)
        vbox.addWidget(self.canvas)
        vbox.addWidget((NavigationToolbar2QT(self.canvas, self)))
        self.setWindowTitle('Picasso: Filter')

    def update_locs(self, locs):
        self.locs = locs
        self.plot()
        self.update()

class HistWindow(PlotWindow):

    def __init__(self, main_window, locs, field):
        self.field = field
        super().__init__(main_window, locs)

    def plot(self):
        # Prepare the data
        data = self.locs[self.field]
        data = data[np.isfinite(data)]
        bins = lib.calculate_optimal_bins(data, 1000)
        # Prepare the figure
        self.figure.clear()
        self.figure.suptitle(self.field)
        axes = self.figure.add_subplot(111)
        axes.hist(data, bins, rwidth=1, linewidth=0)
        data_range = data.ptp()
        axes.set_xlim([bins[0] - 0.05*data_range, data.max() + 0.05*data_range])
        SpanSelector(axes, self.on_span_select, 'horizontal', useblit=True, rectprops=dict(facecolor='green', alpha=0.2))
        self.canvas.draw()

    def on_span_select(self, xmin, xmax):
        self.locs = self.locs[np.isfinite(self.locs[self.field])]
        self.locs = self.locs[(self.locs[self.field] > xmin) & (self.locs[self.field] < xmax)]
        self.main_window.update_locs(self.locs)
        self.main_window.log_filter(self.field, xmin.item(), xmax.item())
        self.plot()

    def closeEvent(self, event):
        self.main_window.hist_windows[self.field] = None
        event.accept()

class Hist2DWindow(PlotWindow):

    def __init__(self, main_window, locs, field_x, field_y):
        self.field_x = field_x
        self.field_y = field_y
        super().__init__(main_window, locs)

    def plot(self):
        # Prepare the data
        x = self.locs[self.field_x]
        y = self.locs[self.field_y]
        valid = (np.isfinite(x) & np.isfinite(y))
        x = x[valid]
        y = y[valid]
        # Prepare the figure
        self.figure.clear()
        axes = self.figure.add_subplot(111)
        # Start hist2 version
        bins_x = lib.calculate_optimal_bins(x, 1000)
        bins_y = lib.calculate_optimal_bins(y, 1000)
        counts, x_edges, y_edges, image = axes.hist2d(x, y, bins=[bins_x, bins_y], norm=LogNorm())
        x_range = x.ptp()
        axes.set_xlim([bins_x[0] - 0.05*x_range, x.max() + 0.05*x_range])
        y_range = y.ptp()
        axes.set_ylim([bins_y[0] - 0.05*y_range, y.max() + 0.05*y_range])
        self.figure.colorbar(image, ax=axes)
        axes.grid(False)
        axes.get_xaxis().set_label_text(self.field_x)
        axes.get_yaxis().set_label_text(self.field_y)
        self.selector = RectangleSelector(axes, self.on_rect_select, useblit=True, rectprops=dict(facecolor='green',
                                                                                alpha=0.2,
                                                                                fill=True))
        self.canvas.draw()

    def on_rect_select(self, press_event, release_event):
        x1, y1 = press_event.xdata, press_event.ydata
        x2, y2 = release_event.xdata, release_event.ydata
        xmin = min(x1, x2)
        xmax = max(x1, x2)
        ymin = min(y1, y2)
        ymax = max(y1, y2)
        self.locs = self.locs[np.isfinite(self.locs[self.field_x])]
        self.locs = self.locs[np.isfinite(self.locs[self.field_y])]
        self.locs = self.locs[(self.locs[self.field_x] > xmin) & (self.locs[self.field_x] < xmax)]
        self.locs = self.locs[(self.locs[self.field_y] > ymin) & (self.locs[self.field_y] < ymax)]
        self.main_window.update_locs(self.locs)
        self.main_window.log_filter(self.field_x, xmin.item(), xmax.item())
        self.main_window.log_filter(self.field_y, ymin.item(), ymax.item())
        self.plot()

    def closeEvent(self, event):
        self.main_window.hist2d_windows[self.field_x][self.field_y] = None
        event.accept()

class Window(QtGui.QMainWindow):

    def __init__(self):
        super().__init__()
        # Init GUI
        self.setWindowTitle('Picasso: Filter')
        self.resize(1100, 750)
        this_directory = os.path.dirname(os.path.realpath(__file__))
        icon_path = os.path.join(this_directory, 'filter.ico')
        icon = QtGui.QIcon(icon_path)
        self.setWindowIcon(icon)
        menu_bar = self.menuBar()
        file_menu = menu_bar.addMenu('File')
        open_action = file_menu.addAction('Open')
        open_action.setShortcut(QtGui.QKeySequence.Open)
        open_action.triggered.connect(self.open_file_dialog)
        file_menu.addAction(open_action)
        save_action = file_menu.addAction('Save')
        save_action.setShortcut(QtGui.QKeySequence.Save)
        save_action.triggered.connect(self.save_file_dialog)
        file_menu.addAction(save_action)
        plot_menu = menu_bar.addMenu('Plot')
        histogram_action = plot_menu.addAction('Histogram')
        histogram_action.setShortcut('Ctrl+H')
        histogram_action.triggered.connect(self.plot_histogram)
        scatter_action = plot_menu.addAction('2D Histogram')
        scatter_action.setShortcut('Ctrl+D')
        scatter_action.triggered.connect(self.plot_hist2d)
        self.table_view = TableView(self, self)
        self.table_view.setAcceptDrops(True)
        self.setCentralWidget(self.table_view)
        self.hist_windows = {}
        self.hist2d_windows = {}
        self.filter_log = {}
        self.locs = None

    def open_file_dialog(self):
        path = QtGui.QFileDialog.getOpenFileName(self, 'Open localizations', filter='*.hdf5')
        if path:
            self.open(path)

    def open(self, path):
        locs, self.info = io.load_locs(path)
        if self.locs is not None:
            for field in self.locs.dtype.names:
                if self.hist_windows[field]:
                    self.hist_windows[field].close()
                for field_y in self.locs.dtype.names:
                    if self.hist2d_windows[field][field_y]:
                        self.hist_windows[field][field_y].close()
        self.locs_path = path
        self.update_locs(locs)
        for field in self.locs.dtype.names:
            self.hist_windows[field] = None
            self.hist2d_windows[field] = {}
            for field_y in self.locs.dtype.names:
                self.hist2d_windows[field][field_y] = None
            self.filter_log[field] = None

    def plot_histogram(self):
        selection_model = self.table_view.selectionModel()
        indices = selection_model.selectedColumns()
        if len(indices) > 0:
            for index in indices:
                index = index.column()
                field = self.locs.dtype.names[index]
                if not self.hist_windows[field]:
                    self.hist_windows[field] = HistWindow(self, self.locs, field)
                self.hist_windows[field].show()

    def plot_hist2d(self):
        selection_model = self.table_view.selectionModel()
        indices = selection_model.selectedColumns()
        if len(indices) == 2:
            indices = [index.column() for index in indices]
            field_x, field_y = [self.locs.dtype.names[index] for index in indices]
            if not self.hist2d_windows[field_x][field_y]:
                self.hist2d_windows[field_x][field_y] = Hist2DWindow(self, self.locs, field_x, field_y)
            self.hist2d_windows[field_x][field_y].show()

    def update_locs(self, locs):
        self.locs = locs
        table_model = TableModel(self.locs, self)
        self.table_view.setModel(table_model)
        for field, hist_window in self.hist_windows.items():
            if hist_window:
                hist_window.update_locs(locs)
        for field_x, hist2d_windows in self.hist2d_windows.items():
            for field_y, hist2d_window in hist2d_windows.items():
                if hist2d_window:
                    hist2d_window.update_locs(locs)

    def log_filter(self, field, xmin, xmax):
        if self.filter_log[field]:
            self.filter_log[field][0] = max(xmin, self.filter_log[field][0])
            self.filter_log[field][1] = min(xmax, self.filter_log[field][1])
        else:
            self.filter_log[field] = [xmin, xmax]

    def save_file_dialog(self):
        base, ext = os.path.splitext(self.locs_path)
        out_path = base + '_filter.hdf5'
        path = QtGui.QFileDialog.getSaveFileName(self, 'Save localizations', out_path, filter='*.hdf5')
        if path:
            filter_info = self.filter_log.copy()
            filter_info.update({'Generated by': 'Picasso Filter'})
            info = self.info + [filter_info]
            io.save_locs(path, self.locs, info)

    def closeEvent(self, event):
        QtGui.qApp.closeAllWindows()

if __name__ == '__main__':
    app = QtGui.QApplication(sys.argv)
    window = Window()
    window.show()

    def excepthook(type, value, tback):
        message = ''.join(traceback.format_exception(type, value, tback))
        errorbox = QtGui.QMessageBox.critical(window, 'An error occured', message)
        errorbox.exec_()
        sys.__excepthook__(type, value, tback)
    sys.excepthook = excepthook

    sys.exit(app.exec_())
<commit_msg>fixed a bug that prevented redrawing axis labels in 2d hists. optimization in the tableview<commit_after>"""
    gui/filter
    ~~~~~~~~~~~~~~~~~~~~

    Graphical user interface for filtering localization lists

    :author: Joerg Schnitzbauer, 2015
"""

import sys
import traceback
from PyQt4 import QtCore, QtGui
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg, NavigationToolbar2QT
import matplotlib.pyplot as plt
from matplotlib.widgets import SpanSelector, RectangleSelector
from matplotlib.colors import LogNorm
import numpy as np
import os.path
from picasso import io, lib

plt.style.use('ggplot')

class TableModel(QtCore.QAbstractTableModel):

    def __init__(self, locs, parent=None):
        super().__init__(parent)
        self.locs = locs
        try:
            self._column_count = len(locs[0])
        except IndexError:
            self._column_count = 0
        self._row_count = self.locs.shape[0]

    def columnCount(self, parent):
        return self._column_count

    def rowCount(self, parent):
        return self._row_count

    def data(self, index, role=QtCore.Qt.DisplayRole):
        if role == QtCore.Qt.DisplayRole:
            data = self.locs[index.row()][index.column()]
            return str(data)
        return None

    def headerData(self, section, orientation, role):
        if role == QtCore.Qt.DisplayRole:
            if orientation == QtCore.Qt.Horizontal:
                return self.locs.dtype.names[section]
            elif orientation == QtCore.Qt.Vertical:
                return section
        return None

class TableView(QtGui.QTableView):

    def __init__(self, window, parent=None):
        super().__init__(parent)
        self.window = window

    def dragEnterEvent(self, event):
        if event.mimeData().hasUrls():
            event.accept()
        else:
            event.ignore()

    def dragMoveEvent(self, event):
        event.accept()

    def dropEvent(self, event):
        urls = event.mimeData().urls()
        path = urls[0].toLocalFile()
        extension = os.path.splitext(path)[1].lower()
        if extension == '.hdf5':
            self.window.open(path)

class PlotWindow(QtGui.QWidget):

    def __init__(self, main_window, locs):
        super().__init__()
        self.main_window = main_window
        self.locs = locs
        self.figure = plt.Figure()
        self.canvas = FigureCanvasQTAgg(self.figure)
        self.plot()
        vbox = QtGui.QVBoxLayout()
        self.setLayout(vbox)
        vbox.addWidget(self.canvas)
        vbox.addWidget((NavigationToolbar2QT(self.canvas, self)))
        self.setWindowTitle('Picasso: Filter')

    def update_locs(self, locs):
        self.locs = locs
        self.plot()
        self.update()

class HistWindow(PlotWindow):

    def __init__(self, main_window, locs, field):
        self.field = field
        super().__init__(main_window, locs)

    def plot(self):
        # Prepare the data
        data = self.locs[self.field]
        data = data[np.isfinite(data)]
        bins = lib.calculate_optimal_bins(data, 1000)
        # Prepare the figure
        self.figure.clear()
        self.figure.suptitle(self.field)
        axes = self.figure.add_subplot(111)
        axes.hist(data, bins, rwidth=1, linewidth=0)
        data_range = data.ptp()
        axes.set_xlim([bins[0] - 0.05*data_range, data.max() + 0.05*data_range])
        self.span = SpanSelector(axes, self.on_span_select, 'horizontal', useblit=True, rectprops=dict(facecolor='green', alpha=0.2))
        self.canvas.draw()

    def on_span_select(self, xmin, xmax):
        print('span selected')
        self.locs = self.locs[np.isfinite(self.locs[self.field])]
        self.locs = self.locs[(self.locs[self.field] > xmin) & (self.locs[self.field] < xmax)]
        self.main_window.update_locs(self.locs)
        self.main_window.log_filter(self.field, xmin.item(), xmax.item())
        self.plot()

    def closeEvent(self, event):
        self.main_window.hist_windows[self.field] = None
        event.accept()

class Hist2DWindow(PlotWindow):

    def __init__(self, main_window, locs, field_x, field_y):
        self.field_x = field_x
        self.field_y = field_y
        super().__init__(main_window, locs)
        self.resize(1000, 800)

    def plot(self):
        # Prepare the data
        x = self.locs[self.field_x]
        y = self.locs[self.field_y]
        valid = (np.isfinite(x) & np.isfinite(y))
        x = x[valid]
        y = y[valid]
        # Prepare the figure
        self.figure.clear()
        # self.canvas.figure = self.figure
        axes = self.figure.add_subplot(111)
        # Start hist2 version
        bins_x = lib.calculate_optimal_bins(x, 1000)
        bins_y = lib.calculate_optimal_bins(y, 1000)
        counts, x_edges, y_edges, image = axes.hist2d(x, y, bins=[bins_x, bins_y], norm=LogNorm())
        x_range = x.ptp()
        axes.set_xlim([bins_x[0] - 0.05*x_range, x.max() + 0.05*x_range])
        y_range = y.ptp()
        axes.set_ylim([bins_y[0] - 0.05*y_range, y.max() + 0.05*y_range])
        self.figure.colorbar(image, ax=axes)
        axes.grid(False)
        axes.get_xaxis().set_label_text(self.field_x)
        axes.get_yaxis().set_label_text(self.field_y)
        self.selector = RectangleSelector(axes, self.on_rect_select, useblit=False,
                                          rectprops=dict(facecolor='green', alpha=0.2, fill=True))
        self.canvas.draw()

    def on_rect_select(self, press_event, release_event):
        x1, y1 = press_event.xdata, press_event.ydata
        x2, y2 = release_event.xdata, release_event.ydata
        xmin = min(x1, x2)
        xmax = max(x1, x2)
        ymin = min(y1, y2)
        ymax = max(y1, y2)
        self.locs = self.locs[np.isfinite(self.locs[self.field_x])]
        self.locs = self.locs[np.isfinite(self.locs[self.field_y])]
        self.locs = self.locs[(self.locs[self.field_x] > xmin) & (self.locs[self.field_x] < xmax)]
        self.locs = self.locs[(self.locs[self.field_y] > ymin) & (self.locs[self.field_y] < ymax)]
        self.main_window.update_locs(self.locs)
        self.main_window.log_filter(self.field_x, xmin.item(), xmax.item())
        self.main_window.log_filter(self.field_y, ymin.item(), ymax.item())
        self.plot()

    def closeEvent(self, event):
        self.main_window.hist2d_windows[self.field_x][self.field_y] = None
        event.accept()

class Window(QtGui.QMainWindow):

    def __init__(self):
        super().__init__()
        # Init GUI
        self.setWindowTitle('Picasso: Filter')
        self.resize(1100, 750)
        this_directory = os.path.dirname(os.path.realpath(__file__))
        icon_path = os.path.join(this_directory, 'filter.ico')
        icon = QtGui.QIcon(icon_path)
        self.setWindowIcon(icon)
        menu_bar = self.menuBar()
        file_menu = menu_bar.addMenu('File')
        open_action = file_menu.addAction('Open')
        open_action.setShortcut(QtGui.QKeySequence.Open)
        open_action.triggered.connect(self.open_file_dialog)
        file_menu.addAction(open_action)
        save_action = file_menu.addAction('Save')
        save_action.setShortcut(QtGui.QKeySequence.Save)
        save_action.triggered.connect(self.save_file_dialog)
        file_menu.addAction(save_action)
        plot_menu = menu_bar.addMenu('Plot')
        histogram_action = plot_menu.addAction('Histogram')
        histogram_action.setShortcut('Ctrl+H')
        histogram_action.triggered.connect(self.plot_histogram)
        scatter_action = plot_menu.addAction('2D Histogram')
        scatter_action.setShortcut('Ctrl+D')
        scatter_action.triggered.connect(self.plot_hist2d)
        self.table_view = TableView(self, self)
        self.table_view.setAcceptDrops(True)
        self.setCentralWidget(self.table_view)
        self.hist_windows = {}
        self.hist2d_windows = {}
        self.filter_log = {}
        self.locs = None

    def open_file_dialog(self):
        path = QtGui.QFileDialog.getOpenFileName(self, 'Open localizations', filter='*.hdf5')
        if path:
            self.open(path)

    def open(self, path):
        locs, self.info = io.load_locs(path)
        if self.locs is not None:
            for field in self.locs.dtype.names:
                if self.hist_windows[field]:
                    self.hist_windows[field].close()
                for field_y in self.locs.dtype.names:
                    if self.hist2d_windows[field][field_y]:
                        self.hist_windows[field][field_y].close()
        self.locs_path = path
        self.update_locs(locs)
        for field in self.locs.dtype.names:
            self.hist_windows[field] = None
            self.hist2d_windows[field] = {}
            for field_y in self.locs.dtype.names:
                self.hist2d_windows[field][field_y] = None
            self.filter_log[field] = None

    def plot_histogram(self):
        selection_model = self.table_view.selectionModel()
        indices = selection_model.selectedColumns()
        if len(indices) > 0:
            for index in indices:
                index = index.column()
                field = self.locs.dtype.names[index]
                if not self.hist_windows[field]:
                    self.hist_windows[field] = HistWindow(self, self.locs, field)
                self.hist_windows[field].show()

    def plot_hist2d(self):
        selection_model = self.table_view.selectionModel()
        indices = selection_model.selectedColumns()
        if len(indices) == 2:
            indices = [index.column() for index in indices]
            field_x, field_y = [self.locs.dtype.names[index] for index in indices]
            if not self.hist2d_windows[field_x][field_y]:
                self.hist2d_windows[field_x][field_y] = Hist2DWindow(self, self.locs, field_x, field_y)
            self.hist2d_windows[field_x][field_y].show()

    def update_locs(self, locs):
        self.locs = locs
        table_model = TableModel(self.locs, self)
        self.table_view.setModel(table_model)
        for field, hist_window in self.hist_windows.items():
            if hist_window:
                hist_window.update_locs(locs)
        for field_x, hist2d_windows in self.hist2d_windows.items():
            for field_y, hist2d_window in hist2d_windows.items():
                if hist2d_window:
                    hist2d_window.update_locs(locs)

    def log_filter(self, field, xmin, xmax):
        if self.filter_log[field]:
            self.filter_log[field][0] = max(xmin, self.filter_log[field][0])
            self.filter_log[field][1] = min(xmax, self.filter_log[field][1])
        else:
            self.filter_log[field] = [xmin, xmax]

    def save_file_dialog(self):
        base, ext = os.path.splitext(self.locs_path)
        out_path = base + '_filter.hdf5'
        path = QtGui.QFileDialog.getSaveFileName(self, 'Save localizations', out_path, filter='*.hdf5')
        if path:
            filter_info = self.filter_log.copy()
            filter_info.update({'Generated by': 'Picasso Filter'})
            info = self.info + [filter_info]
            io.save_locs(path, self.locs, info)

    def closeEvent(self, event):
        QtGui.qApp.closeAllWindows()

if __name__ == '__main__':
    app = QtGui.QApplication(sys.argv)
    window = Window()
    window.show()

    def excepthook(type, value, tback):
        message = ''.join(traceback.format_exception(type, value, tback))
        errorbox = QtGui.QMessageBox.critical(window, 'An error occured', message)
        errorbox.exec_()
        sys.__excepthook__(type, value, tback)
    sys.excepthook = excepthook

    sys.exit(app.exec_())
<|endoftext|>


<commit_before>from feedly.utils import chunks
from feedly.tasks import fanout_operation
from feedly.tasks import follow_many
from feedly.feeds.base import UserBaseFeed

# functions used in tasks need to be at the main level of the module
def add_operation(feed, activities, batch_interface):
    feed.add_many(activities, batch_interface=batch_interface)

def remove_operation(feed, activities, batch_interface):
    feed.remove_many(activities, batch_interface=batch_interface)

class BaseFeedly(object):
    pass

class Feedly(BaseFeedly):
    follow_activity_limit = 5000
    fanout_chunk_size = 1000

    feed_classes = dict()
    user_feed_class = UserBaseFeed

    def __init__(self):
        '''
        This manager is built specifically for the love feed

        :feed_class the feed
        :user_feed_class where user activity gets stored (defaults to same as :feed_class param)

        '''
        pass

    def get_feeds(self, user_id):
        '''
        get the feed that contains the sum of all activity
        from feeds :user_id is subscribed to

        '''
        return dict([(k, feed(user_id)) for k, feed in self.feed_classes.items()])

    def get_user_feed(self, user_id):
        '''
        feed where activity from :user_id is saved

        '''
        return self.user_feed_class(user_id)

    def add_user_activity(self, user_id, activity):
        '''
        Store the new activity and then fanout to user followers

        '''
        self.get_user_feed(user_id).insert_activity(activity)

        user_feed = self.get_user_feed(user_id)
        user_feed.add(activity)
        self._fanout(
            self.feed_classes,
            user_id,
            add_operation,
            activities=[activity]
        )
        return

    def remove_user_activity(self, user_id, activity):
        '''
        Remove the activity and then fanout to user followers

        '''
        self.feed_class.remove_activity(activity)
        user_feed = self.get_user_feed(user_id)
        user_feed.remove(activity)
        self._fanout(
            self.feed_classes,
            user_id,
            remove_operation,
            activities=[activity]
        )
        return

    def follow_feed(self, feed, target_feed):
        '''
        copies target_feed entries in feed
        '''
        activities = target_feed[:self.follow_activity_limit]
        return feed.add_many(activities)

    def unfollow_feed(self, feed, target_feed):
        '''
        removes entries in target_feed from feed

        '''
        activities = feed[:]  # need to slice
        return feed.remove_many(activities)

    def unfollow_user(self, user_id, target_user_id):
        '''
        user_id stops following target_user_id

        '''
        target_feed = self.get_user_feed(target_user_id)
        for feed in self.get_feeds(user_id):
            self.unfollow_feed(feed, target_feed)

    def follow_user(self, user_id, target_user_id):
        '''
        user_id starts following target_user_id
        '''
        target_feed = self.get_user_feed(target_user_id)
        for user_feed in self.get_feeds(user_id):
            self.follow_feed(user_feed, target_feed)

    def follow_many_users(self, user_id, target_ids, async=True):
        '''
        copies feeds for target_ids in user_id
        :async controls if the operation should be done via celery

        '''
        if async:
            follow_many_fn = follow_many.delay
        else:
            follow_many_fn = follow_many

        follow_many_fn(
            feed=self.get_feed(user_id),
            target_feeds=map(self.get_user_feed, target_ids),
            follow_limit=self.follow_activity_limit
        )

    def get_user_follower_ids(self, user_id):
        '''
        returns the list of ids of user_id followers

        '''
        raise NotImplementedError()

    def _fanout(self, feed_classes, user_id, operation, *args, **kwargs):
        '''
        Generic functionality for running an operation on all of your
        follower's feeds

        It takes the following ids and distributes them per fanout_chunk_size
        '''
        user_ids = self.get_user_follower_ids(user_id)
        user_ids_chunks = chunks(user_ids, self.fanout_chunk_size)
        for ids_chunk in user_ids_chunks:
            fanout_operation.delay(
                self, feed_classes, ids_chunk, operation, *args, **kwargs
            )

    def _fanout_task(self, user_ids, feed_classes, operation, *args, **kwargs):
        '''
        This bit of the fan-out is normally called via an Async task
        this shouldnt do any db queries whatsoever
        '''
        for name, feed_class in feed_classes.items():
            with feed_class.get_timeline_batch_interface() as batch_interface:
                kwargs['batch_interface'] = batch_interface
                for user_id in user_ids:
                    feed = feed_class(user_id)
                    operation(feed, *args, **kwargs)

    def flush(self):
        for name, feed_class in self.feed_classes.items():
            feed_class.flush()
        self.user_feed_class.flush()
<commit_msg>fixed broken user follow<commit_after>from feedly.utils import chunks
from feedly.tasks import fanout_operation
from feedly.tasks import follow_many
from feedly.feeds.base import UserBaseFeed

# functions used in tasks need to be at the main level of the module
def add_operation(feed, activities, batch_interface):
    feed.add_many(activities, batch_interface=batch_interface)

def remove_operation(feed, activities, batch_interface):
    feed.remove_many(activities, batch_interface=batch_interface)

class BaseFeedly(object):
    pass

class Feedly(BaseFeedly):
    follow_activity_limit = 5000
    fanout_chunk_size = 1000

    feed_classes = dict()
    user_feed_class = UserBaseFeed

    def __init__(self):
        '''
        This manager is built specifically for the love feed

        :feed_class the feed
        :user_feed_class where user activity gets stored (defaults to same as :feed_class param)

        '''
        pass

    def get_feeds(self, user_id):
        '''
        get the feed that contains the sum of all activity
        from feeds :user_id is subscribed to

        '''
        return dict([(k, feed(user_id)) for k, feed in self.feed_classes.items()])

    def get_user_feed(self, user_id):
        '''
        feed where activity from :user_id is saved

        '''
        return self.user_feed_class(user_id)

    def add_user_activity(self, user_id, activity):
        '''
        Store the new activity and then fanout to user followers

        '''
        self.get_user_feed(user_id).insert_activity(activity)

        user_feed = self.get_user_feed(user_id)
        user_feed.add(activity)
        self._fanout(
            self.feed_classes,
            user_id,
            add_operation,
            activities=[activity]
        )
        return

    def remove_user_activity(self, user_id, activity):
        '''
        Remove the activity and then fanout to user followers

        '''
        self.feed_class.remove_activity(activity)
        user_feed = self.get_user_feed(user_id)
        user_feed.remove(activity)
        self._fanout(
            self.feed_classes,
            user_id,
            remove_operation,
            activities=[activity]
        )
        return

    def follow_feed(self, feed, target_feed):
        '''
        copies target_feed entries in feed
        '''
        activities = target_feed[:self.follow_activity_limit]
        return feed.add_many(activities)

    def unfollow_feed(self, feed, target_feed):
        '''
        removes entries in target_feed from feed

        '''
        activities = feed[:]  # need to slice
        return feed.remove_many(activities)

    def unfollow_user(self, user_id, target_user_id):
        '''
        user_id stops following target_user_id

        '''
        target_feed = self.get_user_feed(target_user_id)
        for feed in self.get_feeds(user_id):
            self.unfollow_feed(feed, target_feed)

    def follow_user(self, user_id, target_user_id):
        '''
        user_id starts following target_user_id
        '''
        target_feed = self.get_user_feed(target_user_id)
        for feed_name, user_feed in self.get_feeds(user_id).items():
            self.follow_feed(user_feed, target_feed)

    def follow_many_users(self, user_id, target_ids, async=True):
        '''
        copies feeds for target_ids in user_id
        :async controls if the operation should be done via celery

        '''
        if async:
            follow_many_fn = follow_many.delay
        else:
            follow_many_fn = follow_many

        follow_many_fn(
            feed=self.get_feed(user_id),
            target_feeds=map(self.get_user_feed, target_ids),
            follow_limit=self.follow_activity_limit
        )

    def get_user_follower_ids(self, user_id):
        '''
        returns the list of ids of user_id followers

        '''
        raise NotImplementedError()

    def _fanout(self, feed_classes, user_id, operation, *args, **kwargs):
        '''
        Generic functionality for running an operation on all of your
        follower's feeds

        It takes the following ids and distributes them per fanout_chunk_size
        '''
        user_ids = self.get_user_follower_ids(user_id)
        user_ids_chunks = chunks(user_ids, self.fanout_chunk_size)
        for ids_chunk in user_ids_chunks:
            fanout_operation.delay(
                self, feed_classes, ids_chunk, operation, *args, **kwargs
            )

    def _fanout_task(self, user_ids, feed_classes, operation, *args, **kwargs):
        '''
        This bit of the fan-out is normally called via an Async task
        this shouldnt do any db queries whatsoever
        '''
        for name, feed_class in feed_classes.items():
            with feed_class.get_timeline_batch_interface() as batch_interface:
                kwargs['batch_interface'] = batch_interface
                for user_id in user_ids:
                    feed = feed_class(user_id)
                    operation(feed, *args, **kwargs)

    def flush(self):
        for name, feed_class in self.feed_classes.items():
            feed_class.flush()
        self.user_feed_class.flush()
<|endoftext|>


<commit_before>import discord
from discord.ext import commands
from .utils import checks
from cogs.utils.dataIO import dataIO, fileIO
from __main__ import send_cmd_help

import json
import os
import asyncio
import aiohttp
import datetime

try: # check if BeautifulSoup4 is installed
	from bs4 import BeautifulSoup
	soupAvailable = True
except:
	soupAvailable = False

...

class APIError(Exception):
	pass

class APIKeyError(Exception):
	pass

class Gw2tp:
	"""This cog finds tp prices"""

	def __init__(self, bot):
		self.bot = bot
		self.session = aiohttp.ClientSession(loop=self.bot.loop)

	@commands.command(pass_context=True)
	async def tpdata(self, ctx, tpbuyid: str):
		"""This finds the current buy price of an item
		Doesn't require any keys/scopes"""
		user = ctx.message.author
		try:
			commerce = 'commerce/prices/'
			endpoint = commerce + tpbuyid
			results = await self.call_api(endpoint)
		except APIKeyError as e:
			await self.bot.say(e)
			return
		except APIError as e:
			await self.bot.say("{0.mention}, API has responded with the following error: "
							   "`{1}`".format(user, e))
			return
		buyprice = results["buys"]["unit_price"]
		sellprice = results ["sells"]["unit_price"]
		buyprice = self.gold_to_coins(buyprice)
		sellprice = self.gold_to_coins(sellprice)
		data = discord.Embed(description=None)
		data.add_field(name="Buy price", value=buyprice)
		data.add_field(name="Sell price", value=sellprice)

		try:
			await self.bot.say(embed=data)
		except discord.HTTPException:
			await self.bot.say("Need permission to embed links")

	async def call_api(self, endpoint):
		apiserv = 'https://api.guildwars2.com/v2/'
		url = apiserv + endpoint
		async with self.session.get(url) as r:
			results = await r.json()
		if "error" in results:
			raise APIError("The API is dead!")
		if "text" in results:
			raise APIError(results["text"])
		return results

def check_folders():
	if not os.path.exists("data/guildwars2"):
		print("Creating data/guildwars2")
		os.makedirs("data/guildwars2")

def check_files():
	files = {
		"gamedata.json": {},
		"settings.json": {"ENABLED": False},
		"language.json": {},
		"keys.json": {},
		"build.json": {"id": None}  # Yay legacy support
	}

	for filename, value in files.items():
		if not os.path.isfile("data/guildwars2/{}".format(filename)):
			print("Creating empty {}".format(filename))
			dataIO.save_json("data/guildwars2/{}".format(filename), value)

def gold_to_coins(self, money):
	gold, remainder = divmod(money, 10000)
	silver, copper = divmod(remainder, 100)
	if not gold:
		if not silver:
			return "{0} copper".format(copper)
		else:
			return "{0} silver and {1} copper".format(silver, copper)
	else:
		return "{0} gold, {1} silver and {2} copper".format(gold, silver, copper)

def setup(bot):
	if soupAvailable:
		bot.add_cog(Gw2tp(bot))
	else:
		raise RuntimeError("You need to run `pip3 install beautifulsoup4`")
	check_folders()
	check_files()<commit_msg>moved gold to coins to correct place<commit_after>import discord
from discord.ext import commands
from .utils import checks
from cogs.utils.dataIO import dataIO, fileIO
from __main__ import send_cmd_help

import json
import os
import asyncio
import aiohttp
import datetime

try: # check if BeautifulSoup4 is installed
	from bs4 import BeautifulSoup
	soupAvailable = True
except:
	soupAvailable = False

...

class APIError(Exception):
	pass

class APIKeyError(Exception):
	pass

class Gw2tp:
	"""This cog finds tp prices"""

	def __init__(self, bot):
		self.bot = bot
		self.session = aiohttp.ClientSession(loop=self.bot.loop)

	@commands.command(pass_context=True)
	async def tpdata(self, ctx, tpbuyid: str):
		"""This finds the current buy price of an item
		Doesn't require any keys/scopes"""
		user = ctx.message.author
		try:
			commerce = 'commerce/prices/'
			endpoint = commerce + tpbuyid
			results = await self.call_api(endpoint)
		except APIKeyError as e:
			await self.bot.say(e)
			return
		except APIError as e:
			await self.bot.say("{0.mention}, API has responded with the following error: "
							   "`{1}`".format(user, e))
			return
		buyprice = results["buys"]["unit_price"]
		sellprice = results ["sells"]["unit_price"]
		buyprice = self.gold_to_coins(buyprice)
		sellprice = self.gold_to_coins(sellprice)
		data = discord.Embed(description=None)
		data.add_field(name="Buy price", value=buyprice)
		data.add_field(name="Sell price", value=sellprice)

		try:
			await self.bot.say(embed=data)
		except discord.HTTPException:
			await self.bot.say("Need permission to embed links")

	async def call_api(self, endpoint):
		apiserv = 'https://api.guildwars2.com/v2/'
		url = apiserv + endpoint
		async with self.session.get(url) as r:
			results = await r.json()
		if "error" in results:
			raise APIError("The API is dead!")
		if "text" in results:
			raise APIError(results["text"])
		return results

	def gold_to_coins(self, money):
	gold, remainder = divmod(money, 10000)
	silver, copper = divmod(remainder, 100)
	if not gold:
		if not silver:
			return "{0} copper".format(copper)
		else:
			return "{0} silver and {1} copper".format(silver, copper)
	else:
		return "{0} gold, {1} silver and {2} copper".format(gold, silver, copper)

def check_folders():
	if not os.path.exists("data/guildwars2"):
		print("Creating data/guildwars2")
		os.makedirs("data/guildwars2")

def check_files():
	files = {
		"gamedata.json": {},
		"settings.json": {"ENABLED": False},
		"language.json": {},
		"keys.json": {},
		"build.json": {"id": None}  # Yay legacy support
	}

	for filename, value in files.items():
		if not os.path.isfile("data/guildwars2/{}".format(filename)):
			print("Creating empty {}".format(filename))
			dataIO.save_json("data/guildwars2/{}".format(filename), value)

def setup(bot):
	if soupAvailable:
		bot.add_cog(Gw2tp(bot))
	else:
		raise RuntimeError("You need to run `pip3 install beautifulsoup4`")
	check_folders()
	check_files()<|endoftext|>


<commit_before># ------------------------------------------------------------------------
# coding=utf-8
# ------------------------------------------------------------------------

import json
import logging

from django.conf import settings as django_settings
from django.contrib import admin
from django.contrib.admin.views import main
from django.db.models import Q
from django.http import (HttpResponse, HttpResponseBadRequest,
    HttpResponseForbidden, HttpResponseNotFound, HttpResponseServerError)
from django.utils.html import escape
from django.utils.safestring import mark_safe
from django.utils.translation import ugettext_lazy as _, ugettext

from mptt.exceptions import InvalidMove
from mptt.forms import MPTTAdminForm

from feincms import settings
from feincms.extensions import ExtensionModelAdmin

# ------------------------------------------------------------------------
def django_boolean_icon(field_val, alt_text=None, title=None):
    """
    Return HTML code for a nice representation of true/false.
    """

    # Origin: contrib/admin/templatetags/admin_list.py
    BOOLEAN_MAPPING = { True: 'yes', False: 'no', None: 'unknown' }
    alt_text = alt_text or BOOLEAN_MAPPING[field_val]
    if title is not None:
        title = 'title="%s" ' % title
    else:
        title = ''
    return mark_safe(u'' %
            (django_settings.STATIC_URL, BOOLEAN_MAPPING[field_val], alt_text, title))

def _build_tree_structure(cls):
    """
    Build an in-memory representation of the item tree, trying to keep
    database accesses down to a minimum. The returned dictionary looks like
    this (as json dump):

        {"6": [7, 8, 10]
         "7": [12],
         "8": [],
         ...
         }
    """
    all_nodes = { }

    mptt_opts = cls._mptt_meta

    for p_id, parent_id in cls.objects.order_by(mptt_opts.tree_id_attr, mptt_opts.left_attr).values_list("pk", "%s_id" % mptt_opts.parent_attr):
        all_nodes[p_id] = []

        if parent_id:
            if not all_nodes.has_key(parent_id):
                # This happens very rarely, but protect against parents that
                # we have yet to iteratove over.
                all_nodes[parent_id] = []
            all_nodes[parent_id].append(p_id)

    return all_nodes

# ------------------------------------------------------------------------
def ajax_editable_boolean_cell(item, attr, text='', override=None):
    """
    Generate a html snippet for showing a boolean value on the admin page.
    Item is an object, attr is the attribute name we should display. Text
    is an optional explanatory text to be included in the output.

    This function will emit code to produce a checkbox input with its state
    corresponding to the item.attr attribute if no override value is passed.
    This input is wired to run a JS ajax updater to toggle the value.

    If override is passed in, ignores the attr attribute and returns a
    static image for the override boolean with no user interaction possible
    (useful for "disabled and you can't change it" situations).
    """
    if text:
        text = ' (%s)' % unicode(text)

    if override is not None:
        a = [ django_boolean_icon(override, text), text ]
    else:
        value = getattr(item, attr)
        a = [
              '<input type="checkbox"',
              value and ' checked="checked"' or '',
              ' onclick="inplace_toggle_boolean(%d, \'%s\').then($.fn.recolorRows);"' % (item.pk, attr),
              ' />',
              text,
            ]

    a.insert(0, '
' % ( attr, item.pk ))
    a.append('
')
    return unicode(''.join(a))

# ------------------------------------------------------------------------
def ajax_editable_boolean(attr, short_description):
    """
    Convenience function: Assign the return value of this method to a variable
    of your ModelAdmin class and put the variable name into list_display.

    Example::

        class MyTreeEditor(TreeEditor):
            list_display = ('__unicode__', 'active_toggle')

            active_toggle = ajax_editable_boolean('active', _('is active'))
    """
    def _fn(self, item):
        return ajax_editable_boolean_cell(item, attr)
    _fn.allow_tags = True
    _fn.short_description = short_description
    _fn.editable_boolean_field = attr
    return _fn

# ------------------------------------------------------------------------
class ChangeList(main.ChangeList):
    """
    Custom ``ChangeList`` class which ensures that the tree entries are always
    ordered in depth-first order (order by ``tree_id``, ``lft``).
    """

    def __init__(self, request, *args, **kwargs):
        self.user = request.user
        super(ChangeList, self).__init__(request, *args, **kwargs)

    def get_query_set(self, *args, **kwargs):
        mptt_opts = self.model._mptt_meta
        return super(ChangeList, self).get_query_set(*args, **kwargs).order_by(mptt_opts.tree_id_attr, mptt_opts.left_attr)

    def get_results(self, request):
        mptt_opts = self.model._mptt_meta
        if settings.FEINCMS_TREE_EDITOR_INCLUDE_ANCESTORS:
            clauses = [Q(**{
                mptt_opts.tree_id_attr: tree_id,
                mptt_opts.left_attr + '__lte': lft,
                mptt_opts.right_attr + '__gte': rght,
                }
                ) for lft, rght, tree_id in \
                    self.query_set.values_list(mptt_opts.left_attr, mptt_opts.right_attr, mptt_opts.tree_id_attr)]
            if clauses:
                self.query_set = self.model._default_manager.filter(reduce(lambda p, q: p|q, clauses))

        super(ChangeList, self).get_results(request)

        for item in self.result_list:
            if settings.FEINCMS_TREE_EDITOR_OBJECT_PERMISSIONS:
                item.feincms_editable = self.model_admin.has_change_permission(request, item)
            else:
                item.feincms_editable = True

# ------------------------------------------------------------------------
# MARK: -
# ------------------------------------------------------------------------

class TreeEditor(ExtensionModelAdmin):
    """
    The ``TreeEditor`` modifies the standard Django administration change list
    to a drag-drop enabled interface for django-mptt_-managed Django models.

    .. _django-mptt: https://github.com/django-mptt/django-mptt/
    """

    form = MPTTAdminForm

    if settings.FEINCMS_TREE_EDITOR_INCLUDE_ANCESTORS:
        # Make sure that no pagination is displayed. Slicing is disabled anyway,
        # therefore this value does not have an influence on the queryset
        list_per_page = 999999999

    def __init__(self, *args, **kwargs):
        super(TreeEditor, self).__init__(*args, **kwargs)

        self.list_display = list(self.list_display)

        if 'indented_short_title' not in self.list_display:
            if self.list_display[0] == 'action_checkbox':
                self.list_display[1] = 'indented_short_title'
            else:
                self.list_display[0] = 'indented_short_title'
        self.list_display_links = ('indented_short_title',)

        opts = self.model._meta
        self.change_list_template = [
            'admin/feincms/%s/%s/tree_editor.html' % (opts.app_label, opts.object_name.lower()),
            'admin/feincms/%s/tree_editor.html' % opts.app_label,
            'admin/feincms/tree_editor.html',
            ]

    def editable(self, item):
        return getattr(item, 'feincms_editable', True)

    def indented_short_title(self, item):
        """
        Generate a short title for an object, indent it depending on
        the object's depth in the hierarchy.
        """
        mptt_opts = item._mptt_meta
        r = ''
        try:
            url = item.get_absolute_url()
        except (AttributeError,):
            url = None

        if url:
            r = '<input type="hidden" class="medialibrary_file_path" value="%s" id="_refkey_%d" />' % (
                        url,
                        item.pk
                      )

        editable_class = ''
        if not getattr(item, 'feincms_editable', True):
            editable_class = ' tree-item-not-editable'

        r += '  ' % (
                item.pk, editable_class, 14+getattr(item, mptt_opts.level_attr)*18)
#        r += ''
        if hasattr(item, 'short_title'):
            r += escape(item.short_title())
        else:
            r += escape(unicode(item))
#        r += ''
        return mark_safe(r)
    indented_short_title.short_description = _('title')
    indented_short_title.allow_tags = True

    def _collect_editable_booleans(self):
        """
        Collect all fields marked as editable booleans. We do not
        want the user to be able to edit arbitrary fields by crafting
        an AJAX request by hand.
        """
        if hasattr(self, '_ajax_editable_booleans'):
            return

        self._ajax_editable_booleans = {}

        for field in self.list_display:
            # The ajax_editable_boolean return value has to be assigned
            # to the ModelAdmin class
            try:
                item = getattr(self.__class__, field)
            except (AttributeError, TypeError):
                continue

            attr = getattr(item, 'editable_boolean_field', None)
            if attr:
                if hasattr(item, 'editable_boolean_result'):
                    result_func = item.editable_boolean_result
                else:
                    def _fn(attr):
                        return lambda self, instance: [ajax_editable_boolean_cell(instance, attr)]
                    result_func = _fn(attr)
                self._ajax_editable_booleans[attr] = result_func

    def _refresh_changelist_caches(self):
        """
        Refresh information used to show the changelist tree structure such as
        inherited active/inactive states etc.

        XXX: This is somewhat hacky, but since it's an internal method, so be it.
        """

        pass

    def _toggle_boolean(self, request):
        """
        Handle an AJAX toggle_boolean request
        """
        try:
            item_id = int(request.POST.get('item_id', None))
            attr = str(request.POST.get('attr', None))
        except:
            return HttpResponseBadRequest("Malformed request")

        if not request.user.is_staff:
            logging.warning("Denied AJAX request by non-staff %s to toggle boolean %s for object #%s", request.user, attr, item_id)
            return HttpResponseForbidden("You do not have permission to access this object")

        self._collect_editable_booleans()

        if not attr in self._ajax_editable_booleans:
            return HttpResponseBadRequest("not a valid attribute %s" % attr)

        try:
            obj = self.model._default_manager.get(pk=item_id)
        except self.model.DoesNotExist:
            return HttpResponseNotFound("Object does not exist")

        if not self.has_change_permission(request, obj=obj):
            logging.warning("Denied AJAX request by %s to toggle boolean %s for object %s", request.user, attr, item_id)
            return HttpResponseForbidden("You do not have permission to access this object")

        new_state = not getattr(obj, attr)
        logging.info("Processing request by %s to toggle %s on #%d %s to %s",
                     request.user, attr, obj.pk, obj, "on" if new_state else "off")

        try:
            before_data = self._ajax_editable_booleans[attr](self, obj)

            setattr(obj, attr, new_state)
            obj.save()

            self._refresh_changelist_caches() # ???: Perhaps better a post_save signal?

            # Construct html snippets to send back to client for status update
            data = self._ajax_editable_booleans[attr](self, obj)

        except Exception:
            logging.exception("Unhandled exception while toggling %s on %s", attr, obj)
            return HttpResponseServerError("Unable to toggle %s on %s" % (attr, obj))

        # Weed out unchanged cells to keep the updates small. This assumes
        # that the order a possible get_descendents() returns does not change
        # before and after toggling this attribute. Unlikely, but still...
        return HttpResponse(json.dumps(
            [b for a, b in zip(before_data, data) if a != b]
            ), content_type="application/json")

    def get_changelist(self, request, **kwargs):
        return ChangeList

    def changelist_view(self, request, extra_context=None, *args, **kwargs):
        """
        Handle the changelist view, the django view for the model instances
        change list/actions page.
        """

        if 'actions_column' not in self.list_display:
            self.list_display.append('actions_column')

        # handle common AJAX requests
        if request.is_ajax():
            cmd = request.POST.get('__cmd')
            if cmd == 'toggle_boolean':
                return self._toggle_boolean(request)
            elif cmd == 'move_node':
                return self._move_node(request)
            else:
                return HttpResponseBadRequest('Oops. AJAX request not understood.')

        self._refresh_changelist_caches()

        extra_context = extra_context or {}
        extra_context['tree_structure'] = mark_safe(json.dumps(
                                                    _build_tree_structure(self.model)))

        return super(TreeEditor, self).changelist_view(request, extra_context, *args, **kwargs)

    def has_change_permission(self, request, obj=None):
        """
        Implement a lookup for object level permissions. Basically the same as
        ModelAdmin.has_change_permission, but also passes the obj parameter in.
        """
        if settings.FEINCMS_TREE_EDITOR_OBJECT_PERMISSIONS:
            opts = self.opts
            r = request.user.has_perm(opts.app_label + '.' + opts.get_change_permission(), obj)
        else:
            r = True

        return r and super(TreeEditor, self).has_change_permission(request, obj)

    def has_delete_permission(self, request, obj=None):
        """
        Implement a lookup for object level permissions. Basically the same as
        ModelAdmin.has_delete_permission, but also passes the obj parameter in.
        """
        if settings.FEINCMS_TREE_EDITOR_OBJECT_PERMISSIONS:
            opts = self.opts
            r = request.user.has_perm(opts.app_label + '.' + opts.get_delete_permission(), obj)
        else:
            r = True

        return r and super(TreeEditor, self).has_delete_permission(request, obj)

    def _move_node(self, request):
        if hasattr(self.model.objects, 'move_node'):
            tree_manager = self.model.objects
        else:
            tree_manager = self.model._tree_manager

        cut_item = tree_manager.get(pk=request.POST.get('cut_item'))
        pasted_on = tree_manager.get(pk=request.POST.get('pasted_on'))
        position = request.POST.get('position')

        if position in ('last-child', 'left', 'right'):
            try:
                tree_manager.move_node(cut_item, pasted_on, position)
            except InvalidMove, e:
                self.message_user(request, unicode(e))
                return HttpResponse('FAIL')

            # Ensure that model save has been run
            cut_item = self.model.objects.get(pk=cut_item.pk)
            cut_item.save()

            self.message_user(request, ugettext('%s has been moved to a new position.') %
                cut_item)
            return HttpResponse('OK')

        self.message_user(request, ugettext('Did not understand moving instruction.'))
        return HttpResponse('FAIL')

    def _actions_column(self, instance):
        return ['

',]

    def actions_column(self, instance):
        return u' '.join(self._actions_column(instance))
    actions_column.allow_tags = True
    actions_column.short_description = _('actions')

    def delete_selected_tree(self, modeladmin, request, queryset):
        """
        Deletes multiple instances and makes sure the MPTT fields get
        recalculated properly.  (Because merely doing a bulk delete doesn't
        trigger the post_delete hooks.)
        """
        n = 0
        for obj in queryset:
            obj.delete()
            n += 1
        self.message_user(request, _("Successfully deleted %s items.") % n)

    def get_actions(self, request):
        actions = super(TreeEditor, self).get_actions(request)
        if 'delete_selected' in actions:
            actions['delete_selected'] = (
                self.delete_selected_tree,
                'delete_selected',
                _("Delete selected %(verbose_name_plural)s"))
        return actions
<commit_msg>Check that short_title is really callable<commit_after># ------------------------------------------------------------------------
# coding=utf-8
# ------------------------------------------------------------------------

import json
import logging

from django.conf import settings as django_settings
from django.contrib import admin
from django.contrib.admin.views import main
from django.db.models import Q
from django.http import (HttpResponse, HttpResponseBadRequest,
    HttpResponseForbidden, HttpResponseNotFound, HttpResponseServerError)
from django.utils.html import escape
from django.utils.safestring import mark_safe
from django.utils.translation import ugettext_lazy as _, ugettext

from mptt.exceptions import InvalidMove
from mptt.forms import MPTTAdminForm

from feincms import settings
from feincms.extensions import ExtensionModelAdmin

# ------------------------------------------------------------------------
def django_boolean_icon(field_val, alt_text=None, title=None):
    """
    Return HTML code for a nice representation of true/false.
    """

    # Origin: contrib/admin/templatetags/admin_list.py
    BOOLEAN_MAPPING = { True: 'yes', False: 'no', None: 'unknown' }
    alt_text = alt_text or BOOLEAN_MAPPING[field_val]
    if title is not None:
        title = 'title="%s" ' % title
    else:
        title = ''
    return mark_safe(u'' %
            (django_settings.STATIC_URL, BOOLEAN_MAPPING[field_val], alt_text, title))

def _build_tree_structure(cls):
    """
    Build an in-memory representation of the item tree, trying to keep
    database accesses down to a minimum. The returned dictionary looks like
    this (as json dump):

        {"6": [7, 8, 10]
         "7": [12],
         "8": [],
         ...
         }
    """
    all_nodes = { }

    mptt_opts = cls._mptt_meta

    for p_id, parent_id in cls.objects.order_by(mptt_opts.tree_id_attr, mptt_opts.left_attr).values_list("pk", "%s_id" % mptt_opts.parent_attr):
        all_nodes[p_id] = []

        if parent_id:
            if not all_nodes.has_key(parent_id):
                # This happens very rarely, but protect against parents that
                # we have yet to iteratove over.
                all_nodes[parent_id] = []
            all_nodes[parent_id].append(p_id)

    return all_nodes

# ------------------------------------------------------------------------
def ajax_editable_boolean_cell(item, attr, text='', override=None):
    """
    Generate a html snippet for showing a boolean value on the admin page.
    Item is an object, attr is the attribute name we should display. Text
    is an optional explanatory text to be included in the output.

    This function will emit code to produce a checkbox input with its state
    corresponding to the item.attr attribute if no override value is passed.
    This input is wired to run a JS ajax updater to toggle the value.

    If override is passed in, ignores the attr attribute and returns a
    static image for the override boolean with no user interaction possible
    (useful for "disabled and you can't change it" situations).
    """
    if text:
        text = ' (%s)' % unicode(text)

    if override is not None:
        a = [ django_boolean_icon(override, text), text ]
    else:
        value = getattr(item, attr)
        a = [
              '<input type="checkbox"',
              value and ' checked="checked"' or '',
              ' onclick="inplace_toggle_boolean(%d, \'%s\').then($.fn.recolorRows);"' % (item.pk, attr),
              ' />',
              text,
            ]

    a.insert(0, '
' % ( attr, item.pk ))
    a.append('
')
    return unicode(''.join(a))

# ------------------------------------------------------------------------
def ajax_editable_boolean(attr, short_description):
    """
    Convenience function: Assign the return value of this method to a variable
    of your ModelAdmin class and put the variable name into list_display.

    Example::

        class MyTreeEditor(TreeEditor):
            list_display = ('__unicode__', 'active_toggle')

            active_toggle = ajax_editable_boolean('active', _('is active'))
    """
    def _fn(self, item):
        return ajax_editable_boolean_cell(item, attr)
    _fn.allow_tags = True
    _fn.short_description = short_description
    _fn.editable_boolean_field = attr
    return _fn

# ------------------------------------------------------------------------
class ChangeList(main.ChangeList):
    """
    Custom ``ChangeList`` class which ensures that the tree entries are always
    ordered in depth-first order (order by ``tree_id``, ``lft``).
    """

    def __init__(self, request, *args, **kwargs):
        self.user = request.user
        super(ChangeList, self).__init__(request, *args, **kwargs)

    def get_query_set(self, *args, **kwargs):
        mptt_opts = self.model._mptt_meta
        return super(ChangeList, self).get_query_set(*args, **kwargs).order_by(mptt_opts.tree_id_attr, mptt_opts.left_attr)

    def get_results(self, request):
        mptt_opts = self.model._mptt_meta
        if settings.FEINCMS_TREE_EDITOR_INCLUDE_ANCESTORS:
            clauses = [Q(**{
                mptt_opts.tree_id_attr: tree_id,
                mptt_opts.left_attr + '__lte': lft,
                mptt_opts.right_attr + '__gte': rght,
                }
                ) for lft, rght, tree_id in \
                    self.query_set.values_list(mptt_opts.left_attr, mptt_opts.right_attr, mptt_opts.tree_id_attr)]
            if clauses:
                self.query_set = self.model._default_manager.filter(reduce(lambda p, q: p|q, clauses))

        super(ChangeList, self).get_results(request)

        for item in self.result_list:
            if settings.FEINCMS_TREE_EDITOR_OBJECT_PERMISSIONS:
                item.feincms_editable = self.model_admin.has_change_permission(request, item)
            else:
                item.feincms_editable = True

# ------------------------------------------------------------------------
# MARK: -
# ------------------------------------------------------------------------

class TreeEditor(ExtensionModelAdmin):
    """
    The ``TreeEditor`` modifies the standard Django administration change list
    to a drag-drop enabled interface for django-mptt_-managed Django models.

    .. _django-mptt: https://github.com/django-mptt/django-mptt/
    """

    form = MPTTAdminForm

    if settings.FEINCMS_TREE_EDITOR_INCLUDE_ANCESTORS:
        # Make sure that no pagination is displayed. Slicing is disabled anyway,
        # therefore this value does not have an influence on the queryset
        list_per_page = 999999999

    def __init__(self, *args, **kwargs):
        super(TreeEditor, self).__init__(*args, **kwargs)

        self.list_display = list(self.list_display)

        if 'indented_short_title' not in self.list_display:
            if self.list_display[0] == 'action_checkbox':
                self.list_display[1] = 'indented_short_title'
            else:
                self.list_display[0] = 'indented_short_title'
        self.list_display_links = ('indented_short_title',)

        opts = self.model._meta
        self.change_list_template = [
            'admin/feincms/%s/%s/tree_editor.html' % (opts.app_label, opts.object_name.lower()),
            'admin/feincms/%s/tree_editor.html' % opts.app_label,
            'admin/feincms/tree_editor.html',
            ]

    def editable(self, item):
        return getattr(item, 'feincms_editable', True)

    def indented_short_title(self, item):
        """
        Generate a short title for an object, indent it depending on
        the object's depth in the hierarchy.
        """
        mptt_opts = item._mptt_meta
        r = ''
        try:
            url = item.get_absolute_url()
        except (AttributeError,):
            url = None

        if url:
            r = '<input type="hidden" class="medialibrary_file_path" value="%s" id="_refkey_%d" />' % (
                        url,
                        item.pk
                      )

        editable_class = ''
        if not getattr(item, 'feincms_editable', True):
            editable_class = ' tree-item-not-editable'

        r += '  ' % (
                item.pk, editable_class, 14+getattr(item, mptt_opts.level_attr)*18)
#        r += ''
        if hasattr(item, 'short_title') and callable(item.short_title):
            r += escape(item.short_title())
        else:
            r += escape(unicode(item))
#        r += ''
        return mark_safe(r)
    indented_short_title.short_description = _('title')
    indented_short_title.allow_tags = True

    def _collect_editable_booleans(self):
        """
        Collect all fields marked as editable booleans. We do not
        want the user to be able to edit arbitrary fields by crafting
        an AJAX request by hand.
        """
        if hasattr(self, '_ajax_editable_booleans'):
            return

        self._ajax_editable_booleans = {}

        for field in self.list_display:
            # The ajax_editable_boolean return value has to be assigned
            # to the ModelAdmin class
            try:
                item = getattr(self.__class__, field)
            except (AttributeError, TypeError):
                continue

            attr = getattr(item, 'editable_boolean_field', None)
            if attr:
                if hasattr(item, 'editable_boolean_result'):
                    result_func = item.editable_boolean_result
                else:
                    def _fn(attr):
                        return lambda self, instance: [ajax_editable_boolean_cell(instance, attr)]
                    result_func = _fn(attr)
                self._ajax_editable_booleans[attr] = result_func

    def _refresh_changelist_caches(self):
        """
        Refresh information used to show the changelist tree structure such as
        inherited active/inactive states etc.

        XXX: This is somewhat hacky, but since it's an internal method, so be it.
        """

        pass

    def _toggle_boolean(self, request):
        """
        Handle an AJAX toggle_boolean request
        """
        try:
            item_id = int(request.POST.get('item_id', None))
            attr = str(request.POST.get('attr', None))
        except:
            return HttpResponseBadRequest("Malformed request")

        if not request.user.is_staff:
            logging.warning("Denied AJAX request by non-staff %s to toggle boolean %s for object #%s", request.user, attr, item_id)
            return HttpResponseForbidden("You do not have permission to access this object")

        self._collect_editable_booleans()

        if not attr in self._ajax_editable_booleans:
            return HttpResponseBadRequest("not a valid attribute %s" % attr)

        try:
            obj = self.model._default_manager.get(pk=item_id)
        except self.model.DoesNotExist:
            return HttpResponseNotFound("Object does not exist")

        if not self.has_change_permission(request, obj=obj):
            logging.warning("Denied AJAX request by %s to toggle boolean %s for object %s", request.user, attr, item_id)
            return HttpResponseForbidden("You do not have permission to access this object")

        new_state = not getattr(obj, attr)
        logging.info("Processing request by %s to toggle %s on #%d %s to %s",
                     request.user, attr, obj.pk, obj, "on" if new_state else "off")

        try:
            before_data = self._ajax_editable_booleans[attr](self, obj)

            setattr(obj, attr, new_state)
            obj.save()

            self._refresh_changelist_caches() # ???: Perhaps better a post_save signal?

            # Construct html snippets to send back to client for status update
            data = self._ajax_editable_booleans[attr](self, obj)

        except Exception:
            logging.exception("Unhandled exception while toggling %s on %s", attr, obj)
            return HttpResponseServerError("Unable to toggle %s on %s" % (attr, obj))

        # Weed out unchanged cells to keep the updates small. This assumes
        # that the order a possible get_descendents() returns does not change
        # before and after toggling this attribute. Unlikely, but still...
        return HttpResponse(json.dumps(
            [b for a, b in zip(before_data, data) if a != b]
            ), content_type="application/json")

    def get_changelist(self, request, **kwargs):
        return ChangeList

    def changelist_view(self, request, extra_context=None, *args, **kwargs):
        """
        Handle the changelist view, the django view for the model instances
        change list/actions page.
        """

        if 'actions_column' not in self.list_display:
            self.list_display.append('actions_column')

        # handle common AJAX requests
        if request.is_ajax():
            cmd = request.POST.get('__cmd')
            if cmd == 'toggle_boolean':
                return self._toggle_boolean(request)
            elif cmd == 'move_node':
                return self._move_node(request)
            else:
                return HttpResponseBadRequest('Oops. AJAX request not understood.')

        self._refresh_changelist_caches()

        extra_context = extra_context or {}
        extra_context['tree_structure'] = mark_safe(json.dumps(
                                                    _build_tree_structure(self.model)))

        return super(TreeEditor, self).changelist_view(request, extra_context, *args, **kwargs)

    def has_change_permission(self, request, obj=None):
        """
        Implement a lookup for object level permissions. Basically the same as
        ModelAdmin.has_change_permission, but also passes the obj parameter in.
        """
        if settings.FEINCMS_TREE_EDITOR_OBJECT_PERMISSIONS:
            opts = self.opts
            r = request.user.has_perm(opts.app_label + '.' + opts.get_change_permission(), obj)
        else:
            r = True

        return r and super(TreeEditor, self).has_change_permission(request, obj)

    def has_delete_permission(self, request, obj=None):
        """
        Implement a lookup for object level permissions. Basically the same as
        ModelAdmin.has_delete_permission, but also passes the obj parameter in.
        """
        if settings.FEINCMS_TREE_EDITOR_OBJECT_PERMISSIONS:
            opts = self.opts
            r = request.user.has_perm(opts.app_label + '.' + opts.get_delete_permission(), obj)
        else:
            r = True

        return r and super(TreeEditor, self).has_delete_permission(request, obj)

    def _move_node(self, request):
        if hasattr(self.model.objects, 'move_node'):
            tree_manager = self.model.objects
        else:
            tree_manager = self.model._tree_manager

        cut_item = tree_manager.get(pk=request.POST.get('cut_item'))
        pasted_on = tree_manager.get(pk=request.POST.get('pasted_on'))
        position = request.POST.get('position')

        if position in ('last-child', 'left', 'right'):
            try:
                tree_manager.move_node(cut_item, pasted_on, position)
            except InvalidMove, e:
                self.message_user(request, unicode(e))
                return HttpResponse('FAIL')

            # Ensure that model save has been run
            cut_item = self.model.objects.get(pk=cut_item.pk)
            cut_item.save()

            self.message_user(request, ugettext('%s has been moved to a new position.') %
                cut_item)
            return HttpResponse('OK')

        self.message_user(request, ugettext('Did not understand moving instruction.'))
        return HttpResponse('FAIL')

    def _actions_column(self, instance):
        return ['

',]

    def actions_column(self, instance):
        return u' '.join(self._actions_column(instance))
    actions_column.allow_tags = True
    actions_column.short_description = _('actions')

    def delete_selected_tree(self, modeladmin, request, queryset):
        """
        Deletes multiple instances and makes sure the MPTT fields get
        recalculated properly.  (Because merely doing a bulk delete doesn't
        trigger the post_delete hooks.)
        """
        n = 0
        for obj in queryset:
            obj.delete()
            n += 1
        self.message_user(request, _("Successfully deleted %s items.") % n)

    def get_actions(self, request):
        actions = super(TreeEditor, self).get_actions(request)
        if 'delete_selected' in actions:
            actions['delete_selected'] = (
                self.delete_selected_tree,
                'delete_selected',
                _("Delete selected %(verbose_name_plural)s"))
        return actions
<|endoftext|>


<commit_before>import re
import sys
import logging

import Evtx.Evtx
import Evtx.Nodes
import Evtx.Views

import evtxtract.utils
import evtxtract.templates

logger = logging.getLogger(__name__)

# via: https://github.com/vivisect/synapse/blob/master/synapse/compat.py
major = sys.version_info.major
minor = sys.version_info.minor
micro = sys.version_info.micro

majmin = (major,minor)
version = (major,minor,micro)

if version < (3,0,0):
    def isstr(s):
        return type(s) in (str,unicode)
else:
    def isstr(s):
        return type(s) == str

class Template(object):
    substitition_re = re.compile("\[(Conditional|Normal) Substitution\(index=(\d+), type=(\d+)\)\]")

    def __init__(self, eid, xml):
        self.eid = eid
        self.xml = xml

        self._cached_placeholders = None
        self._cached_id = None

    def get_id(self):
        """
        @rtype: str
        @return: A string that can be parsed into constraints describing what
          types of subsitutions this template can accept.
          Short example: 1100-[0|4|c]-[1|4|c]-[2|6|c]-[3|6|c]
        """
        if self._cached_id is not None:
            return self._cached_id

        ret = [str(self.eid)]
        for index, type_, mode in self._get_placeholders():
            if mode:
                mode_str = "c"
            else:
                mode_str = "n"
            ret.append("[%s|%s|%s]" % (index, type_, mode_str))

        self._cached_id = "-".join(ret)
        return self._cached_id

    def _get_placeholders(self):
        """
        Get descriptors for each of the substitutions required by this
          template.

        Tuple schema: (index, type, is_conditional)

        @rtype: list of (int, int, boolean)
        """
        if self._cached_placeholders is not None:
            return self._cached_placeholders

        ret = []
        for mode, index, type_ in Template.substitition_re.findall(self.xml):
            ret.append((int(index), int(type_), mode == "Conditional"))

        self._cached_placeholders = sorted(ret, key=lambda p: p[0])
        return self._cached_placeholders

    def match_substitutions(self, substitutions):
        """
        Checks to see if the provided set of substitutions match the
          placeholder values required by this template.

        Note, this is only a best guess.  The number of substitutions
          *may* be greater than the number of available slots. So we
          must only check the slot and substitution types.

        @type substitutions: list of (int, str)
        @param substitutions: Tuple schema (type, value)
        @rtype: boolean
        """
        logger = logging.getLogger("match_substitutions")
        placeholders = self._get_placeholders()
        logger.debug("Substitutions: %s", str(substitutions))
        logger.debug("Constraints: %s", str(placeholders))
        if len(placeholders) > len(substitutions):
            logger.debug("Failing on lens: %d vs %d",
                         len(placeholders), len(substitutions))
            return False
        if max(placeholders, key=lambda k: k[0])[0] > len(substitutions):
            logger.debug("Failing on max index: %d vs %d",
                         max(placeholders, key=lambda k: k[0])[0],
                         len(substitutions))
            return False

        # it seems that some templates request different values than what are subsequently put in them
        #   specifically, a Hex64 might be put into a SizeType field (EID 4624)
        # this maps from the type described in a template, to possible additional types that a
        #   record can provide for a particular substitution
        overrides = {
            16: set([21])
        }

        for index, type_, is_conditional in placeholders:
            sub_type, sub_value = substitutions[index]
            if is_conditional and sub_type == 0:
                continue
            if sub_type != type_:
                if type_ not in overrides or sub_type not in overrides[type_]:
                    logger.debug("Failing on type comparison, index %d: %d vs %d (mode: %s)",
                                 index, sub_type, type_, is_conditional)
                    return False
                else:
                    logger.debug("Overriding template type %d with substitution type %d", type_, sub_type)
                    continue
        return True

    escape_re = re.compile(r"\\\\(\d)")

    @staticmethod
    def _escape(value):
        """
        Escape the static value to be used in a regular expression
          subsititution. This processes any backreferences and
          makes them plain, escaped sequences.

        @type value: str
        @rtype: str
        """
        return Template.escape_re.sub(r"\\\\\\\\\1", re.escape(value))

    def insert_substitutions(self, substitutions):
        """
        Return a copy of the template with the given substitutions inserted.

        @type substitutions: list of (int, str)
        @param substitutions: an ordered list of (type:int, value:str)
        @rtype: str
        """
        ret = self.xml
        for index, pair in enumerate(substitutions):
            type_, value = pair
            from_pattern = "\[(Normal|Conditional) Substitution\(index=%d, type=\d+\)\]" % index
            if isstr(value):
                value = Template._escape(value)
            else:
                value = str(value)
            ret = re.sub(from_pattern, value, ret)
        return ret

REPLACEMENT_PATTERNS = {
        i: re.compile("\[(Normal|Conditional) Substitution\(index=%d, type=\d+\)\]" % i)
          for i in range(35)}

def make_replacement(template, index, substitution):
    """
    Makes a substitution given a template as a string.

    Implementation is a huge hack that depends on the
    brittle template_format() output.

    @type template: str
    @type index: int
    @type substitution: str
    @rtype: str
    """
    if index not in REPLACEMENT_PATTERNS:
        from_pattern = re.compile("\[(Normal|Conditional) Substitution\(index=%d, type=\d+\)\]" % index)
        REPLACEMENT_PATTERNS[index] = from_pattern
    return REPLACEMENT_PATTERNS[index].sub(substitution, template)

def get_complete_template(root, current_index=0):
    """
    Gets the template from a RootNode while resolving any
    nested templates and fixing up their indices.
    Depth first ordering/indexing.

    Implementation is a huge hack that depends on the
      brittle template_format() output.

    @type root: RootNode
    @type current_index: int
    @rtype: str
    """
    template = Evtx.Views.evtx_template_readable_view(root)  # TODO(wb): make sure this is working

    # walk through each substitution.
    # if its a normal node, continue
    # else its a subtemplate, and we count the number of substitutions _it_ has
    #   so that we can later fixup all the indices
    replacements = []
    for index, substitution in enumerate(root.substitutions()):
        # find all sub-templates
        if not isinstance(substitution, Evtx.Nodes.BXmlTypeNode):
            replacements.append(current_index + index)
            continue
        # TODO(wb): hack here accessing ._root
        subtemplate = get_complete_template(substitution._root,
                                             current_index=current_index + index)
        replacements.append(subtemplate)
        current_index += subtemplate.count("Substitution(index=")
    replacements.reverse()

    # now walk through all the indices and fix them up depth-first
    for i, replacement in enumerate(replacements):
        index = len(replacements) - i - 1
        if isinstance(replacement, int):
            # fixup index
            from_pattern = "index=%d," % index
            to_pattern = "index=%d," % replacement
            template = template.replace(from_pattern, to_pattern)
        if isstr(replacement):
            # insert sub-template
            template = make_replacement(template, index, replacement)
    return template

def get_template(record):
    """
    Given a complete Record, parse out the nodes that make up the Template
      and return it as a Template.

    @type record: Record
    @rtype: Template
    """
    record_xml = Evtx.Views.evtx_record_xml_view(record)
    eid = evtxtract.utils.get_eid(record_xml)
    return Template(eid, get_complete_template(record.root()))
<commit_msg>templates: use six when possible<commit_after>import re
import sys
import logging

import six
import Evtx.Evtx
import Evtx.Nodes
import Evtx.Views

import evtxtract.utils
import evtxtract.templates

logger = logging.getLogger(__name__)

class Template(object):
    substitition_re = re.compile("\[(Conditional|Normal) Substitution\(index=(\d+), type=(\d+)\)\]")

    def __init__(self, eid, xml):
        self.eid = eid
        self.xml = xml

        self._cached_placeholders = None
        self._cached_id = None

    def get_id(self):
        """
        @rtype: str
        @return: A string that can be parsed into constraints describing what
          types of subsitutions this template can accept.
          Short example: 1100-[0|4|c]-[1|4|c]-[2|6|c]-[3|6|c]
        """
        if self._cached_id is not None:
            return self._cached_id

        ret = [str(self.eid)]
        for index, type_, mode in self._get_placeholders():
            if mode:
                mode_str = "c"
            else:
                mode_str = "n"
            ret.append("[%s|%s|%s]" % (index, type_, mode_str))

        self._cached_id = "-".join(ret)
        return self._cached_id

    def _get_placeholders(self):
        """
        Get descriptors for each of the substitutions required by this
          template.

        Tuple schema: (index, type, is_conditional)

        @rtype: list of (int, int, boolean)
        """
        if self._cached_placeholders is not None:
            return self._cached_placeholders

        ret = []
        for mode, index, type_ in Template.substitition_re.findall(self.xml):
            ret.append((int(index), int(type_), mode == "Conditional"))

        self._cached_placeholders = sorted(ret, key=lambda p: p[0])
        return self._cached_placeholders

    def match_substitutions(self, substitutions):
        """
        Checks to see if the provided set of substitutions match the
          placeholder values required by this template.

        Note, this is only a best guess.  The number of substitutions
          *may* be greater than the number of available slots. So we
          must only check the slot and substitution types.

        @type substitutions: list of (int, str)
        @param substitutions: Tuple schema (type, value)
        @rtype: boolean
        """
        logger = logging.getLogger("match_substitutions")
        placeholders = self._get_placeholders()
        logger.debug("Substitutions: %s", str(substitutions))
        logger.debug("Constraints: %s", str(placeholders))
        if len(placeholders) > len(substitutions):
            logger.debug("Failing on lens: %d vs %d",
                         len(placeholders), len(substitutions))
            return False
        if max(placeholders, key=lambda k: k[0])[0] > len(substitutions):
            logger.debug("Failing on max index: %d vs %d",
                         max(placeholders, key=lambda k: k[0])[0],
                         len(substitutions))
            return False

        # it seems that some templates request different values than what are subsequently put in them
        #   specifically, a Hex64 might be put into a SizeType field (EID 4624)
        # this maps from the type described in a template, to possible additional types that a
        #   record can provide for a particular substitution
        overrides = {
            16: set([21])
        }

        for index, type_, is_conditional in placeholders:
            sub_type, sub_value = substitutions[index]
            if is_conditional and sub_type == 0:
                continue
            if sub_type != type_:
                if type_ not in overrides or sub_type not in overrides[type_]:
                    logger.debug("Failing on type comparison, index %d: %d vs %d (mode: %s)",
                                 index, sub_type, type_, is_conditional)
                    return False
                else:
                    logger.debug("Overriding template type %d with substitution type %d", type_, sub_type)
                    continue
        return True

    escape_re = re.compile(r"\\\\(\d)")

    @staticmethod
    def _escape(value):
        """
        Escape the static value to be used in a regular expression
          subsititution. This processes any backreferences and
          makes them plain, escaped sequences.

        @type value: str
        @rtype: str
        """
        return Template.escape_re.sub(r"\\\\\\\\\1", re.escape(value))

    def insert_substitutions(self, substitutions):
        """
        Return a copy of the template with the given substitutions inserted.

        @type substitutions: list of (int, str)
        @param substitutions: an ordered list of (type:int, value:str)
        @rtype: str
        """
        ret = self.xml
        for index, pair in enumerate(substitutions):
            type_, value = pair
            from_pattern = "\[(Normal|Conditional) Substitution\(index=%d, type=\d+\)\]" % index
            if isinstance(value, six.string_types):
                value = Template._escape(value)
            else:
                value = str(value)
            ret = re.sub(from_pattern, value, ret)
        return ret

REPLACEMENT_PATTERNS = {
        i: re.compile("\[(Normal|Conditional) Substitution\(index=%d, type=\d+\)\]" % i)
          for i in range(35)}

def make_replacement(template, index, substitution):
    """
    Makes a substitution given a template as a string.

    Implementation is a huge hack that depends on the
    brittle template_format() output.

    @type template: str
    @type index: int
    @type substitution: str
    @rtype: str
    """
    if index not in REPLACEMENT_PATTERNS:
        from_pattern = re.compile("\[(Normal|Conditional) Substitution\(index=%d, type=\d+\)\]" % index)
        REPLACEMENT_PATTERNS[index] = from_pattern
    return REPLACEMENT_PATTERNS[index].sub(substitution, template)

def get_complete_template(root, current_index=0):
    """
    Gets the template from a RootNode while resolving any
    nested templates and fixing up their indices.
    Depth first ordering/indexing.

    Implementation is a huge hack that depends on the
      brittle template_format() output.

    @type root: RootNode
    @type current_index: int
    @rtype: str
    """
    template = Evtx.Views.evtx_template_readable_view(root)  # TODO(wb): make sure this is working

    # walk through each substitution.
    # if its a normal node, continue
    # else its a subtemplate, and we count the number of substitutions _it_ has
    #   so that we can later fixup all the indices
    replacements = []
    for index, substitution in enumerate(root.substitutions()):
        # find all sub-templates
        if not isinstance(substitution, Evtx.Nodes.BXmlTypeNode):
            replacements.append(current_index + index)
            continue
        # TODO(wb): hack here accessing ._root
        subtemplate = get_complete_template(substitution._root,
                                             current_index=current_index + index)
        replacements.append(subtemplate)
        current_index += subtemplate.count("Substitution(index=")
    replacements.reverse()

    # now walk through all the indices and fix them up depth-first
    for i, replacement in enumerate(replacements):
        index = len(replacements) - i - 1
        if isinstance(replacement, int):
            # fixup index
            from_pattern = "index=%d," % index
            to_pattern = "index=%d," % replacement
            template = template.replace(from_pattern, to_pattern)
        if isinstance(replacement, six.string_types):
            # insert sub-template
            template = make_replacement(template, index, replacement)
    return template

def get_template(record):
    """
    Given a complete Record, parse out the nodes that make up the Template
      and return it as a Template.

    @type record: Record
    @rtype: Template
    """
    record_xml = Evtx.Views.evtx_record_xml_view(record)
    eid = evtxtract.utils.get_eid(record_xml)
    return Template(eid, get_complete_template(record.root()))
<|endoftext|>


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

import os
import sys
from path import path
from charmhelpers.contrib.ansible import AnsibleHooks

# Create the hooks helper, passing a list of hooks which will be
# handled by default by running all sections of the playbook
# tagged with the hook name.

def hook_names(here=path(__file__).parent):
    for name in (x.basename()
                 for x in here.files() if x.islink()):
        yield name

hooks = AnsibleHooks(playbook_path='playbooks/site.yaml',
                     default_hooks=list(hook_names()))

if not os.path.exists('/etc/ansible'):
    os.mkdir('/etc/ansible')
    # Work around some wonkyness with ansible
    with open('/etc/ansible/hosts', 'w') as f:
        f.write('localhost')

if __name__ == "__main__":
    hooks.execute(sys.argv)
<commit_msg>Correcting lint<commit_after>#!/usr/bin/env python

import os
import sys
from path import path
from charmhelpers.contrib.ansible import AnsibleHooks

# Create the hooks helper, passing a list of hooks which will be
# handled by default by running all sections of the playbook
# tagged with the hook name.

def hook_names(here=path(__file__).parent):
    for name in (x.basename()
                 for x in here.files() if x.islink()):
        yield name

hooks = AnsibleHooks(playbook_path='playbooks/site.yaml',
                     default_hooks=list(hook_names()))

if not os.path.exists('/etc/ansible'):
    os.mkdir('/etc/ansible')
    # Work around some wonkyness with ansible
    with open('/etc/ansible/hosts', 'w') as f:
        f.write('localhost')

if __name__ == "__main__":
    hooks.execute(sys.argv)
<|endoftext|>