<commit_before>import qutip.examples
import sys,os,time,syntax
from numpy import arange,floor
from qutip.examples import exconfig
from qutip.examples import *
if os.environ['QUTIP_GUI']=="PYSIDE":
    from PySide import QtGui, QtCore

elif os.environ['QUTIP_GUI']=="PYQT4":
    from PyQt4 import QtGui, QtCore

class Examples(QtGui.QWidget):
    def center(self):
        screen = QtGui.QDesktopWidget().screenGeometry()
        size =  self.frameSize()
        self.move((screen.width()-size.width())/2, (screen.height()-size.height())/2)
    def moveout(self):
        #self.clearFocus()
        for op in arange(0.9,-0.1,-0.1):
            time.sleep(.02)
            self.setWindowOpacity(op)
        #self.setVisible(False)
    def movein(self):
        self.focusWidget()
        self.setVisible(True)
        for op in arange(0.1,1.1,0.1):
            time.sleep(.02)
            self.setWindowOpacity(op)

    def __init__(self,version,direc,parent=None):
        QtGui.QWidget.__init__(self, parent)
        from qutip.examples.examples_text import tab_labels,button_labels,button_desc,button_nums
        #set tab button style
        tab_button_style='QPushButton {font-family: Arial;border-width: 2px;border-color:#666666;border-style: solid;border-radius: 7;background-color: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #B8E2EF, stop: 0.1 #A5DBEB, stop: 0.49 #8CD1E6, stop: 0.5 #7BCAE1, stop: 1 #57BCD9)}'
        #set quit button style
        quit_style='QPushButton {font-family: Arial;border-width: 2px;border-color:#666666;border-style: solid;border-radius: 7;font-size: 16px;background-color: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #FFAAAA, stop: 0.1 #FF9999, stop: 0.49 #FF8888, stop: 0.5 #FF7777, stop: 1 #FF6666)}'

        #WINDOW PROPERTIES
        self.setWindowTitle('QuTiP Examples')
        self.resize(1200, 650)
        self.setMinimumSize(1200, 650)
        self.setMaximumSize(1200, 650)
        self.center()
        self.setFocus()
        self.dir=direc+'/examples/'
        mapper = QtCore.QSignalMapper(self)

        title_font = QtGui.QFont()
        title_font.setFamily("Arial")
        title_font.setBold(True)
        title_font.setPointSize(16)
        title_fm = QtGui.QFontMetrics(title_font)
        #text across top of demos window
        title = QtGui.QLabel(self)
        title.setFont(title_font)
        title_text="Click button once for code preview.  Click again to run example."
        pixelswide = title_fm.width(title_text)
        title.setText(title_text)
        title.setGeometry((self.width()-pixelswide)/2.0, 10, 800, 30)

        font = QtGui.QFont()
        font.setFamily("Arial")
        font.setBold(True)
        font.setPointSize(12)
        fm = QtGui.QFontMetrics(font)

        #QUIT BUTTON-----------------
        quit = HoverExit('Close',self)
        quit.setGeometry(1120, 610, 70, 30)
        quit.setStyleSheet(quit_style)
        quit.setFocusPolicy(QtCore.Qt.NoFocus)
        quit.clicked.connect(self.close)

        #copyright text
        copyright = QtGui.QLabel(self)
        copy_text="Copyright (c) 2011-2012, P. D. Nation & J. R. Johansson"
        copyright.setText(copy_text)
        font.setBold(False)
        copyright.setFont(font)
        copyright.setGeometry(10, 620, 400, 30)
        font.setBold(True)
        #tab widget
        tab_widget = QtGui.QTabWidget()
        tab_widget.move(10,10)
        tab_widget.resize(550,550)
        tab_widget.setFixedSize(550,550)
        #tabs for tab widget
        num_tabs=len(tab_labels)
        tabs=[QtGui.QWidget() for k in range(num_tabs)]
        for k in range(num_tabs):
            tab_widget.addTab(tabs[k],tab_labels[k])
        tab_widget.setTabShape(QtGui.QTabWidget.Triangular)
        tab_widget.setCurrentIndex(exconfig.tab)

        #tab buttons
        self.tab_verts =[QtGui.QVBoxLayout(tabs[k]) for k in xrange(num_tabs)]
        self.num_elems=[len(button_labels[k]) for k in xrange(num_tabs)]
        self.tab_buttons=[[] for j in range(num_tabs)]
        for j in range(num_tabs):
            for k in range(self.num_elems[j]):
                button=HoverButton(button_nums[j][k],button_labels[j][k])
                button.setFont(font)
                if button_nums[j][k]==exconfig.is_green:
                    button.setText('Run Example')
                    button.setStyleSheet('QPushButton {font-family: Arial;border-width: 2px;border-color:#666666;border-style: solid;border-radius: 7;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #E0F1E0, stop: 0.1 #C1E3C1, stop: 0.49 #A3D6A3, stop: 0.5 #84C884, stop: 1 #66BB66)}')
                else:
                    button.setText(button_labels[j][k])
                    button.setStyleSheet(tab_button_style)
                button.setFixedSize(150, 40)
                self.connect(button, QtCore.SIGNAL("clicked()"), mapper, QtCore.SLOT("map()"))
                mapper.setMapping(button,button_nums[j][k])
                self.tab_buttons[j].append(button)

        self.tab_button_desc = [[] for j in range(num_tabs)]
        for j in range(num_tabs):
            for k in range(self.num_elems[j]):
                self.tab_button_desc[j].append(QtGui.QLabel(button_desc[j][k]))

        self.tab_widgets=[[QtGui.QWidget() for k in range(self.num_elems[j])] for j in range(num_tabs)]
        self.tab_horiz_layouts=[[QtGui.QHBoxLayout(self.tab_widgets[j][k]) for k in range(self.num_elems[j])] for j in range(num_tabs)]
        for j in range(num_tabs):
            for k in range(self.num_elems[j]):
                self.tab_horiz_layouts[j][k].addWidget(self.tab_buttons[j][k])
                self.tab_horiz_layouts[j][k].addSpacing(20)
                self.tab_horiz_layouts[j][k].addWidget(self.tab_button_desc[j][k])
        for j in range(num_tabs):
            for k in range(self.num_elems[j]):
                self.tab_verts[j].addWidget(self.tab_widgets[j][k])
            self.tab_verts[j].addStretch()

        #set mapper to on_button_clicked funtions
        self.connect(mapper, QtCore.SIGNAL("mapped(int)"), self.on_button_clicked)
        self.layout = QtGui.QGridLayout(self)
        #create text editor widget
        self.editor = QtGui.QPlainTextEdit()
        self.editor.setReadOnly(True)
        self.editor.resize(550,550)
        self.editor.setFixedSize(550,550)
        #add tabwidget and textwidget to main window
        self.layout.addWidget(tab_widget, 0, 0, 5, 1)
        self.layout.addWidget(self.editor, 0, 3, 5, 1)
        #set text of editor if demos previsouly ran.
        if exconfig.is_green!=0:
            syntax.PythonHighlighter(self.editor.document())
            self.editor.show()
            _text_file=compile("infile = open(self.dir+'ex_"+str(exconfig.is_green)+".py', 'r')",'<string>', 'exec')
            exec(_text_file)
            self.editor.setPlainText(infile.read())

    def on_button_clicked(self,num):
        """
        Receives integers from button click to use for calling example script
        """
        if exconfig.button_num!=num and exconfig.is_green!=num:
            tab_num=int(floor(num/10))-1
            row_num=num % 10
            self.tab_buttons[tab_num][row_num].setStyleSheet('QPushButton {font-family: Arial;border-width: 3px;border-color:#111111;border-style: solid;border-radius: 7;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #E0F1E0, stop: 0.1 #C1E3C1, stop: 0.49 #A3D6A3, stop: 0.5 #84C884, stop: 1 #66BB66)}')
            self.tab_buttons[tab_num][row_num].setText('Run Example')
            exconfig.button_num=num
            old_tab_num=int(floor(exconfig.is_green/10))-1
            old_row_num=exconfig.is_green % 10
            self.tab_buttons[old_tab_num][old_row_num].setStyleSheet('QPushButton {font-family: Arial;border-width: 2px;border-color:#666666;border-style: solid;border-radius: 7;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #B8E2EF, stop: 0.1 #A5DBEB, stop: 0.49 #8CD1E6, stop: 0.5 #7BCAE1, stop: 1 #57BCD9)}')
            self.tab_buttons[old_tab_num][old_row_num].setText(self.tab_buttons[old_tab_num][old_row_num].label)
            exconfig.is_green=num
            syntax.PythonHighlighter(self.editor.document())
            self.editor.show()
            _text_file=compile("infile = open(self.dir+'ex_"+str(exconfig.is_green)+".py', 'r')",'<string>', 'exec')
            exec(_text_file)
            self.editor.setPlainText(infile.read())
        else:
            self.moveout()
            exconfig.tab=int(floor(num/10))-1
            exconfig.option=num
            self.close()

class HoverButton(QtGui.QPushButton):
    def __init__(self,num,label):
            super(HoverButton, self).__init__()
            self.num=num
            self.label=label
    def enterEvent(self,event):
        if exconfig.is_green==self.num:
            self.setStyleSheet('QPushButton {font-family: Arial;border-width: 3px;border-color:#111111;border-style: solid;border-radius: 7;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #E0F1E0, stop: 0.1 #C1E3C1, stop: 0.49 #A3D6A3, stop: 0.5 #84C884, stop: 1 #66BB66)}')
        else:
            self.setStyleSheet('QPushButton {font-family: Arial;border-width: 3px;border-color:#111111;border-style: solid;border-radius: 7;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #B8E2EF, stop: 0.1 #A5DBEB, stop: 0.49 #8CD1E6, stop: 0.5 #7BCAE1, stop: 1 #57BCD9)}')
    def leaveEvent(self,event):
        if exconfig.is_green==self.num:
            self.setStyleSheet('QPushButton {font-family: Arial;border-width: 2px;border-color:#666666;border-style: solid;border-radius: 7;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #E0F1E0, stop: 0.1 #C1E3C1, stop: 0.49 #A3D6A3, stop: 0.5 #84C884, stop: 1 #66BB66)}')
        else:
            self.setStyleSheet('QPushButton {font-family: Arial;border-width: 2px;border-color:#666666;border-style: solid;border-radius: 7;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #B8E2EF, stop: 0.1 #A5DBEB, stop: 0.49 #8CD1E6, stop: 0.5 #7BCAE1, stop: 1 #57BCD9)}')

class HoverExit(QtGui.QPushButton):
    def enterEvent(self,event):
        self.setStyleSheet('QPushButton {font-family: Arial;border-width: 3px;border-color:#111111;border-style: solid;border-radius: 7;font-size: 16px;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #FFAAAA, stop: 0.1 #FF9999, stop: 0.49 #FF8888, stop: 0.5 #FF7777, stop: 1 #FF6666)}')
    def leaveEvent(self,event):
        self.setStyleSheet('QPushButton {font-family: Arial;border-width: 2px;border-color:#666666;border-style: solid;border-radius: 7;font-size: 16px;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #FFAAAA, stop: 0.1 #FF9999, stop: 0.49 #FF8888, stop: 0.5 #FF7777, stop: 1 #FF6666)}')

<commit_msg>examples window enlarged to hold 10 examples per tab<commit_after>import qutip.examples
import sys,os,time,syntax
from numpy import arange,floor
from qutip.examples import exconfig
from qutip.examples import *
if os.environ['QUTIP_GUI']=="PYSIDE":
    from PySide import QtGui, QtCore

elif os.environ['QUTIP_GUI']=="PYQT4":
    from PyQt4 import QtGui, QtCore

class Examples(QtGui.QWidget):
    def center(self):
        screen = QtGui.QDesktopWidget().screenGeometry()
        size =  self.frameSize()
        self.move((screen.width()-size.width())/2, (screen.height()-size.height())/2)
    def moveout(self):
        #self.clearFocus()
        for op in arange(0.9,-0.1,-0.1):
            time.sleep(.02)
            self.setWindowOpacity(op)
        #self.setVisible(False)
    def movein(self):
        self.focusWidget()
        self.setVisible(True)
        for op in arange(0.1,1.1,0.1):
            time.sleep(.02)
            self.setWindowOpacity(op)

    def __init__(self,version,direc,parent=None):
        QtGui.QWidget.__init__(self, parent)
        from qutip.examples.examples_text import tab_labels,button_labels,button_desc,button_nums
        #set tab button style
        tab_button_style='QPushButton {font-family: Arial;border-width: 2px;border-color:#666666;border-style: solid;border-radius: 7;background-color: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #B8E2EF, stop: 0.1 #A5DBEB, stop: 0.49 #8CD1E6, stop: 0.5 #7BCAE1, stop: 1 #57BCD9)}'
        #set quit button style
        quit_style='QPushButton {font-family: Arial;border-width: 2px;border-color:#666666;border-style: solid;border-radius: 7;font-size: 16px;background-color: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #FFAAAA, stop: 0.1 #FF9999, stop: 0.49 #FF8888, stop: 0.5 #FF7777, stop: 1 #FF6666)}'

        #WINDOW PROPERTIES
        self.setWindowTitle('QuTiP Examples')
        self.resize(1200, 775)
        self.setMinimumSize(1200, 775)
        self.setMaximumSize(1200, 775)
        self.center()
        self.setFocus()
        self.dir=direc+'/examples/'
        mapper = QtCore.QSignalMapper(self)

        title_font = QtGui.QFont()
        title_font.setFamily("Arial")
        title_font.setBold(True)
        title_font.setPointSize(16)
        title_fm = QtGui.QFontMetrics(title_font)
        #text across top of demos window
        title = QtGui.QLabel(self)
        title.setFont(title_font)
        title_text="Click button once for code preview.  Click again to run example."
        pixelswide = title_fm.width(title_text)
        title.setText(title_text)
        title.setGeometry((self.width()-pixelswide)/2.0, 10, 800, 30)

        font = QtGui.QFont()
        font.setFamily("Arial")
        font.setBold(True)
        font.setPointSize(12)
        fm = QtGui.QFontMetrics(font)

        #QUIT BUTTON-----------------
        quit = HoverExit('Close',self)
        quit.setGeometry(1120, 730, 70, 40)
        quit.setStyleSheet(quit_style)
        quit.setFocusPolicy(QtCore.Qt.NoFocus)
        quit.clicked.connect(self.close)

        #copyright text
        copyright = QtGui.QLabel(self)
        copy_text="Copyright (c) 2011-2012, P. D. Nation & J. R. Johansson"
        copyright.setText(copy_text)
        font.setBold(False)
        copyright.setFont(font)
        copyright.setGeometry(10, 745, 400, 30)
        font.setBold(True)
        #tab widget
        tab_widget = QtGui.QTabWidget()
        tab_widget.move(10,10)
        tab_widget.resize(550,675)
        tab_widget.setFixedSize(550,675)
        #tabs for tab widget
        num_tabs=len(tab_labels)
        tabs=[QtGui.QWidget() for k in range(num_tabs)]
        for k in range(num_tabs):
            tab_widget.addTab(tabs[k],tab_labels[k])
        tab_widget.setTabShape(QtGui.QTabWidget.Triangular)
        tab_widget.setCurrentIndex(exconfig.tab)

        #tab buttons
        self.tab_verts =[QtGui.QVBoxLayout(tabs[k]) for k in xrange(num_tabs)]
        self.num_elems=[len(button_labels[k]) for k in xrange(num_tabs)]
        self.tab_buttons=[[] for j in range(num_tabs)]
        for j in range(num_tabs):
            for k in range(self.num_elems[j]):
                button=HoverButton(button_nums[j][k],button_labels[j][k])
                button.setFont(font)
                if button_nums[j][k]==exconfig.is_green:
                    button.setText('Run Example')
                    button.setStyleSheet('QPushButton {font-family: Arial;border-width: 2px;border-color:#666666;border-style: solid;border-radius: 7;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #E0F1E0, stop: 0.1 #C1E3C1, stop: 0.49 #A3D6A3, stop: 0.5 #84C884, stop: 1 #66BB66)}')
                else:
                    button.setText(button_labels[j][k])
                    button.setStyleSheet(tab_button_style)
                button.setFixedSize(150, 40)
                self.connect(button, QtCore.SIGNAL("clicked()"), mapper, QtCore.SLOT("map()"))
                mapper.setMapping(button,button_nums[j][k])
                self.tab_buttons[j].append(button)

        self.tab_button_desc = [[] for j in range(num_tabs)]
        for j in range(num_tabs):
            for k in range(self.num_elems[j]):
                self.tab_button_desc[j].append(QtGui.QLabel(button_desc[j][k]))

        self.tab_widgets=[[QtGui.QWidget() for k in range(self.num_elems[j])] for j in range(num_tabs)]
        self.tab_horiz_layouts=[[QtGui.QHBoxLayout(self.tab_widgets[j][k]) for k in range(self.num_elems[j])] for j in range(num_tabs)]
        for j in range(num_tabs):
            for k in range(self.num_elems[j]):
                self.tab_horiz_layouts[j][k].addWidget(self.tab_buttons[j][k])
                self.tab_horiz_layouts[j][k].addSpacing(15)
                self.tab_horiz_layouts[j][k].addWidget(self.tab_button_desc[j][k])
        for j in range(num_tabs):
            for k in range(self.num_elems[j]):
                self.tab_verts[j].addWidget(self.tab_widgets[j][k])
            self.tab_verts[j].addStretch()

        #set mapper to on_button_clicked funtions
        self.connect(mapper, QtCore.SIGNAL("mapped(int)"), self.on_button_clicked)
        self.layout = QtGui.QGridLayout(self)
        #create text editor widget
        self.editor = QtGui.QPlainTextEdit()
        self.editor.setReadOnly(True)
        self.editor.resize(550,675)
        self.editor.setFixedSize(550,675)
        #add tabwidget and textwidget to main window
        self.layout.addWidget(tab_widget, 0, 0, 5, 1)
        self.layout.addWidget(self.editor, 0, 3, 5, 1)
        #set text of editor if demos previsouly ran.
        if exconfig.is_green!=0:
            syntax.PythonHighlighter(self.editor.document())
            self.editor.show()
            _text_file=compile("infile = open(self.dir+'ex_"+str(exconfig.is_green)+".py', 'r')",'<string>', 'exec')
            exec(_text_file)
            self.editor.setPlainText(infile.read())

    def on_button_clicked(self,num):
        """
        Receives integers from button click to use for calling example script
        """
        if exconfig.button_num!=num and exconfig.is_green!=num:
            tab_num=int(floor(num/10))-1
            row_num=num % 10
            self.tab_buttons[tab_num][row_num].setStyleSheet('QPushButton {font-family: Arial;border-width: 3px;border-color:#111111;border-style: solid;border-radius: 7;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #E0F1E0, stop: 0.1 #C1E3C1, stop: 0.49 #A3D6A3, stop: 0.5 #84C884, stop: 1 #66BB66)}')
            self.tab_buttons[tab_num][row_num].setText('Run Example')
            exconfig.button_num=num
            old_tab_num=int(floor(exconfig.is_green/10))-1
            old_row_num=exconfig.is_green % 10
            self.tab_buttons[old_tab_num][old_row_num].setStyleSheet('QPushButton {font-family: Arial;border-width: 2px;border-color:#666666;border-style: solid;border-radius: 7;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #B8E2EF, stop: 0.1 #A5DBEB, stop: 0.49 #8CD1E6, stop: 0.5 #7BCAE1, stop: 1 #57BCD9)}')
            self.tab_buttons[old_tab_num][old_row_num].setText(self.tab_buttons[old_tab_num][old_row_num].label)
            exconfig.is_green=num
            syntax.PythonHighlighter(self.editor.document())
            self.editor.show()
            _text_file=compile("infile = open(self.dir+'ex_"+str(exconfig.is_green)+".py', 'r')",'<string>', 'exec')
            exec(_text_file)
            self.editor.setPlainText(infile.read())
        else:
            self.moveout()
            exconfig.tab=int(floor(num/10))-1
            exconfig.option=num
            self.close()

class HoverButton(QtGui.QPushButton):
    def __init__(self,num,label):
            super(HoverButton, self).__init__()
            self.num=num
            self.label=label
    def enterEvent(self,event):
        if exconfig.is_green==self.num:
            self.setStyleSheet('QPushButton {font-family: Arial;border-width: 3px;border-color:#111111;border-style: solid;border-radius: 7;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #E0F1E0, stop: 0.1 #C1E3C1, stop: 0.49 #A3D6A3, stop: 0.5 #84C884, stop: 1 #66BB66)}')
        else:
            self.setStyleSheet('QPushButton {font-family: Arial;border-width: 3px;border-color:#111111;border-style: solid;border-radius: 7;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #B8E2EF, stop: 0.1 #A5DBEB, stop: 0.49 #8CD1E6, stop: 0.5 #7BCAE1, stop: 1 #57BCD9)}')
    def leaveEvent(self,event):
        if exconfig.is_green==self.num:
            self.setStyleSheet('QPushButton {font-family: Arial;border-width: 2px;border-color:#666666;border-style: solid;border-radius: 7;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #E0F1E0, stop: 0.1 #C1E3C1, stop: 0.49 #A3D6A3, stop: 0.5 #84C884, stop: 1 #66BB66)}')
        else:
            self.setStyleSheet('QPushButton {font-family: Arial;border-width: 2px;border-color:#666666;border-style: solid;border-radius: 7;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #B8E2EF, stop: 0.1 #A5DBEB, stop: 0.49 #8CD1E6, stop: 0.5 #7BCAE1, stop: 1 #57BCD9)}')

class HoverExit(QtGui.QPushButton):
    def enterEvent(self,event):
        self.setStyleSheet('QPushButton {font-family: Arial;border-width: 3px;border-color:#111111;border-style: solid;border-radius: 7;font-size: 16px;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #FFAAAA, stop: 0.1 #FF9999, stop: 0.49 #FF8888, stop: 0.5 #FF7777, stop: 1 #FF6666)}')
    def leaveEvent(self,event):
        self.setStyleSheet('QPushButton {font-family: Arial;border-width: 2px;border-color:#666666;border-style: solid;border-radius: 7;font-size: 16px;background: QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #FFAAAA, stop: 0.1 #FF9999, stop: 0.49 #FF8888, stop: 0.5 #FF7777, stop: 1 #FF6666)}')

<|endoftext|>


<commit_before>from django.http import HttpResponseNotAllowed, HttpResponseForbidden, HttpResponse
from django.core.urlresolvers import reverse
from django.core.cache import cache
from django import get_version as django_version
from decorator import decorator

from datetime import datetime, timedelta

__version__ = '0.2'

def get_version():
    return __version__

def format_error(error):
    return u"Piston/%s (Django %s) crash report:\n\n%s" % \
        (get_version(), django_version(), error)

def create_reply(message, status=200):
    return HttpResponse(message, status=status)

class rc(object):
    """
    Status codes.
    """
    ALL_OK = create_reply('OK', status=200)
    CREATED = create_reply('Created', status=201)
    DELETED = create_reply('', status=204) # 204 says "Don't send a body!"
    BAD_REQUEST = create_reply('Bad Request', status=400)
    FORBIDDEN = create_reply('Forbidden', status=401)
    DUPLICATE_ENTRY = create_reply('Conflict/Duplicate', status=409)
    NOT_HERE = create_reply('Gone', status=410)
    NOT_IMPLEMENTED = create_reply('Not Implemented', status=501)
    THROTTLED = create_reply('Throttled', status=503)

class FormValidationError(Exception):
    def __init__(self, form):
        self.form = form

class HttpStatusCode(Exception):
    def __init__(self, message, code=200):
        self.message = message
        self.code = code

def validate(v_form, operation='POST'):
    @decorator
    def wrap(f, self, request, *a, **kwa):
        form = v_form(getattr(request, operation))

        if form.is_valid():
#            kwa.update({ 'form': form })
            return f(self, request, *a, **kwa)
        else:
            raise FormValidationError(form)
    return wrap

def throttle(max_requests, timeout=60*60, extra=''):
    """
    Simple throttling decorator, caches
    the amount of requests made in cache.

    If used on a view where users are required to
    log in, the username is used, otherwise the
    IP address of the originating request is used.

    Parameters::
     - `max_requests`: The maximum number of requests
     - `timeout`: The timeout for the cache entry (default: 1 hour)
    """
    @decorator
    def wrap(f, self, request, *args, **kwargs):
        if request.user.is_authenticated():
            ident = request.user.username
        else:
            ident = request.META.get('REMOTE_ADDR', None)

        if hasattr(request, 'throttle_extra'):
            """
            Since we want to be able to throttle on a per-
            application basis, it's important that we realize
            that `throttle_extra` might be set on the request
            object. If so, append the identifier name with it.
            """
            ident += ':%s' % str(request.throttle_extra)

        if ident:
            """
            Preferrably we'd use incr/decr here, since they're
            atomic in memcached, but it's in django-trunk so we
            can't use it yet. If someone sees this after it's in
            stable, you can change it here.
            """
            ident += ':%s' % extra

            now = datetime.now()
            ts_key = 'throttle:ts:%s' % ident
            timestamp = cache.get(ts_key)
            offset = now + timedelta(seconds=timeout)

            if timestamp and timestamp < offset:
                t = rc.THROTTLED
                wait = timeout - (offset-timestamp).seconds
                t.content = 'Throttled, wait %d seconds.' % wait

                return t

            count = cache.get(ident, 1)
            cache.set(ident, count+1)

            if count >= max_requests:
                cache.set(ts_key, offset, timeout)
                cache.set(ident, 1)

        return f(self, request, *args, **kwargs)
    return wrap

def coerce_put_post(request):
    if request.method == "PUT":
        request.method = "POST"
        request._load_post_and_files()
        request.method = "PUT"
        request.PUT = request.POST
        del request._post

class Mimer(object):
    TYPES = dict()

    def __init__(self, request):
        self.request = request

    def is_multipart(self):
        return 'Content-Disposition: form-data;' in self.request.raw_post_data

    def loader_for_type(self, ctype):
        """
        Gets a function ref to deserialize content
        for a certain mimetype.
        """
        for loadee, mimes in Mimer.TYPES.iteritems():
            if ctype in mimes:
                return loadee

    def content_type(self):
        return self.request.META.get('CONTENT_TYPE')

    def translate(self):
        """
        Will look at the `Content-type` sent by the client, and maybe
        deserialize the contents into the format they sent. This will
        work for JSON, YAML, XML and Pickle. Since the data is not just
        key-value (and maybe just a list), the data will be placed on
        `request.data` instead, and the handler will have to read from
        there.

        It will also set `request.mimetype` so the handler has an easy
        way to tell what's going on. `request.mimetype` will always be
        None for multipart form data (what your browser sends.)
        """
        ctype = self.content_type()

        if not self.is_multipart() and ctype:
            loadee = self.loader_for_type(ctype)

            try:
                self.request.content_type = ctype
                self.request.data = loadee(self.request.raw_post_data)

                # Reset both POST and PUT from request, as its
                # misleading having their presence around.
                self.request.POST = self.request.PUT = dict()
            except TypeError:
                return rc.BAD_REQUEST # TODO: Handle this in super
            except Exception, e:
                raise

        return self.request

    @classmethod
    def register(cls, loadee, types):
        cls.TYPES[loadee] = types

    @classmethod
    def unregister(cls, loadee):
        return cls.TYPES.pop(loadee)

def translate_mime(request):
    request = Mimer(request).translate()

def require_mime(*mimes):
    """
    Decorator requiring a certain mimetype. There's a nifty
    helper called `require_extended` below which requires everything
    we support except for post-data via form.
    """
    @decorator
    def wrap(f, self, request, *args, **kwargs):
        m = Mimer(request)
        realmimes = set()

        rewrite = { 'json':   'application/json',
                    'yaml':   'application/x-yaml',
                    'xml':    'text/xml',
                    'pickle': 'application/python-pickle' }

        for idx, mime in enumerate(mimes):
            realmimes.add(rewrite.get(mime, mime))

        if not m.content_type() in realmimes:
            return rc.BAD_REQUEST

        return f(self, request, *args, **kwargs)
    return wrap

require_extended = require_mime('json', 'yaml', 'xml', 'pickle')
    <commit_msg>utils.rc.RESPONSE returns a fresh HttpResponse every time now<commit_after>from django.http import HttpResponseNotAllowed, HttpResponseForbidden, HttpResponse
from django.core.urlresolvers import reverse
from django.core.cache import cache
from django import get_version as django_version
from decorator import decorator

from datetime import datetime, timedelta

__version__ = '0.2'

def get_version():
    return __version__

def format_error(error):
    return u"Piston/%s (Django %s) crash report:\n\n%s" % \
        (get_version(), django_version(), error)

class rc_factory(object):
    """
    Status codes.
    """
    CODES = dict(ALL_OK = ('OK', 200),
                 CREATED = ('Created', 201),
                 DELETED = ('', 204), # 204 says "Don't send a body!"
                 BAD_REQUEST = ('Bad Request', 400),
                 FORBIDDEN = ('Forbidden', 401),
                 DUPLICATE_ENTRY = ('Conflict/Duplicate', 409),
                 NOT_HERE = ('Gone', 410),
                 NOT_IMPLEMENTED = ('Not Implemented', 501),
                 THROTTLED = ('Throttled', 503))

    def __getattr__(self, attr):
        """
        Returns a fresh `HttpResponse` when getting
        an "attribute". This is backwards compatible
        with 0.2, which is important.
        """
        try:
            (r, c) = self.CODES.get(attr)
        except TypeError:
            raise AttributeError(attr)

        return HttpResponse(r, c)

rc = rc_factory()

class FormValidationError(Exception):
    def __init__(self, form):
        self.form = form

class HttpStatusCode(Exception):
    def __init__(self, message, code=200):
        self.message = message
        self.code = code

def validate(v_form, operation='POST'):
    @decorator
    def wrap(f, self, request, *a, **kwa):
        form = v_form(getattr(request, operation))

        if form.is_valid():
            return f(self, request, *a, **kwa)
        else:
            raise FormValidationError(form)
    return wrap

def throttle(max_requests, timeout=60*60, extra=''):
    """
    Simple throttling decorator, caches
    the amount of requests made in cache.

    If used on a view where users are required to
    log in, the username is used, otherwise the
    IP address of the originating request is used.

    Parameters::
     - `max_requests`: The maximum number of requests
     - `timeout`: The timeout for the cache entry (default: 1 hour)
    """
    @decorator
    def wrap(f, self, request, *args, **kwargs):
        if request.user.is_authenticated():
            ident = request.user.username
        else:
            ident = request.META.get('REMOTE_ADDR', None)

        if hasattr(request, 'throttle_extra'):
            """
            Since we want to be able to throttle on a per-
            application basis, it's important that we realize
            that `throttle_extra` might be set on the request
            object. If so, append the identifier name with it.
            """
            ident += ':%s' % str(request.throttle_extra)

        if ident:
            """
            Preferrably we'd use incr/decr here, since they're
            atomic in memcached, but it's in django-trunk so we
            can't use it yet. If someone sees this after it's in
            stable, you can change it here.
            """
            ident += ':%s' % extra

            now = datetime.now()
            ts_key = 'throttle:ts:%s' % ident
            timestamp = cache.get(ts_key)
            offset = now + timedelta(seconds=timeout)

            if timestamp and timestamp < offset:
                t = rc.THROTTLED
                wait = timeout - (offset-timestamp).seconds
                t.content = 'Throttled, wait %d seconds.' % wait

                return t

            count = cache.get(ident, 1)
            cache.set(ident, count+1)

            if count >= max_requests:
                cache.set(ts_key, offset, timeout)
                cache.set(ident, 1)

        return f(self, request, *args, **kwargs)
    return wrap

def coerce_put_post(request):
    if request.method == "PUT":
        request.method = "POST"
        request._load_post_and_files()
        request.method = "PUT"
        request.PUT = request.POST
        del request._post

class Mimer(object):
    TYPES = dict()

    def __init__(self, request):
        self.request = request

    def is_multipart(self):
        return 'Content-Disposition: form-data;' in self.request.raw_post_data

    def loader_for_type(self, ctype):
        """
        Gets a function ref to deserialize content
        for a certain mimetype.
        """
        for loadee, mimes in Mimer.TYPES.iteritems():
            if ctype in mimes:
                return loadee

    def content_type(self):
        return self.request.META.get('CONTENT_TYPE')

    def translate(self):
        """
        Will look at the `Content-type` sent by the client, and maybe
        deserialize the contents into the format they sent. This will
        work for JSON, YAML, XML and Pickle. Since the data is not just
        key-value (and maybe just a list), the data will be placed on
        `request.data` instead, and the handler will have to read from
        there.

        It will also set `request.mimetype` so the handler has an easy
        way to tell what's going on. `request.mimetype` will always be
        None for multipart form data (what your browser sends.)
        """
        ctype = self.content_type()

        if not self.is_multipart() and ctype:
            loadee = self.loader_for_type(ctype)

            try:
                self.request.content_type = ctype
                self.request.data = loadee(self.request.raw_post_data)

                # Reset both POST and PUT from request, as its
                # misleading having their presence around.
                self.request.POST = self.request.PUT = dict()
            except TypeError:
                return rc.BAD_REQUEST # TODO: Handle this in super
            except Exception, e:
                raise

        return self.request

    @classmethod
    def register(cls, loadee, types):
        cls.TYPES[loadee] = types

    @classmethod
    def unregister(cls, loadee):
        return cls.TYPES.pop(loadee)

def translate_mime(request):
    request = Mimer(request).translate()

def require_mime(*mimes):
    """
    Decorator requiring a certain mimetype. There's a nifty
    helper called `require_extended` below which requires everything
    we support except for post-data via form.
    """
    @decorator
    def wrap(f, self, request, *args, **kwargs):
        m = Mimer(request)
        realmimes = set()

        rewrite = { 'json':   'application/json',
                    'yaml':   'application/x-yaml',
                    'xml':    'text/xml',
                    'pickle': 'application/python-pickle' }

        for idx, mime in enumerate(mimes):
            realmimes.add(rewrite.get(mime, mime))

        if not m.content_type() in realmimes:
            return rc.BAD_REQUEST

        return f(self, request, *args, **kwargs)
    return wrap

require_extended = require_mime('json', 'yaml', 'xml', 'pickle')
    <|endoftext|>


<commit_before># -*- encoding: utf-8 -*-
'''
HubbleStack Nova plugin for auditing installed packages.

This module checks all of a system's local packages and reports if the package
is vulnerable to a known cve. The cve vunlerablities are gathered via the url in
the yaml profile, and that data cached at the path
/var/cache/salt/minion/cve_scan_cache/<os_name>_<version>.json

:maintainer: HubbleStack / jaredhanson11
:maturity: 2016.7.0
:platform: Linux
:requires: SaltStack

This audit module requires yaml data to execute. It will search the local
directory for any .yaml files, and if it finds a top-level 'cve_scan_v2' key, it
will use that data.

Sample YAML data with inline comments:

cve_scan_v2:
    # Seconds until the local cache expires
    ttl: 86400
    # Source of cve data
    url: http://vulners.com/
    # Optional control tag
    control:
        # minimum score, vulnerabilities with a smaller
        #   score added to 'Controlled' output
        score: 3

The source of the cve data can be http://vulners.com/, salt://path/to/json, and
any other url that returns cve data in json format. If the url contains
vulners.com, then this module will use the local system's os and os version to
dynamically query vulner.com/api/v3 for cve data specifically related to your
system. If the url doesn't contain vulners.com, it will query the exact url, so
that endpoint must return cve data specific to the system you are scanning.

The cve data json must be formatted as follows:

[

{'_source': {'affectedPackage': [{'OS': 'CentOS',
                                    'OSVersion': '7',
                                    'operator': 'lt',
                                    'packageFilename': 'krb5-server-1.13.2-12.el7_2.x86_64.rpm',
                                    'packageName': 'krb5-server',
                                    'packageVersion': '1.13.2-12.el7_2'},
                                   {'OS': 'CentOS',
                                    'OSVersion': '7',
                                    'operator': 'lt',
                                    'packageFilename': 'krb5-libs-1.13.2-12.el7_2.i686.rpm',
                                    'packageName': 'krb5-libs',
                                    'packageVersion': '1.13.2-12.el7_2'}
                                    ]
              'cvelist': ['CVE-2015-8631',
                           'CVE-2015-8630',
                           'CVE-2015-8629'],
              'cvss': {'score': 6.8}
              'href': 'http://lists.centos.org/pipermail/centos-announce/2016-March/021788.html',
              'reporter': 'CentOS Project',
              'title': 'Moderate krb5 Security Update'
            }
    },
...

 ]

'''
from __future__ import absolute_import
import logging

import fnmatch
import hashlib
import json
import os
import re
import requests

from distutils.version import LooseVersion
from time import time as current_time
from zipfile import ZipFile

import salt
import salt.utils

log = logging.getLogger(__name__)

def __virtual__():
    return not salt.utils.is_windows()

def audit(data_list, tags, debug=False, **kwargs):
    '''
    Main audit function. See module docstring for more information on usage.
    '''
    os_version = __grains__.get('osmajorrelease', None)
    if os_version is None:
        os_version = __grains__.get('osrelease', None)
    os_name = __grains__['os'].lower()

    log.debug("os_version: %s, os_name: %s", os_version, os_name)

    endpoints = []

    # Go through yaml to check for cve_scan_v2,
    #    if its present, check for a cached version
    #    of the scan.
    for profile, data in data_list:

        if 'cve_scan_v2' in data:

            ttl = data['cve_scan_v2']['ttl']
            url = data['cve_scan_v2']['url']
            control = data['cve_scan_v2'].get('control', {})
            # Ability to add more controls easily, in control dict
            min_score = float(control.get('score', 0))
            urlhash = hashlib.md5(url).hexdigest()
            cached_json = os.path.join(__opts__['cachedir'],
                                       'cve_scan_cache',
                                       '%s.json' % urlhash)
            cached_zip = os.path.join(__opts__['cachedir'],
                                      'cve_scan_cache',
                                      '%s.zip' % urlhash)
            # Make cache directory and all parent directories if it doesn't exist.
            if not os.path.exists(os.path.dirname(cached_json)):
                os.makedirs(os.path.dirname(cached_json))
            cache = _get_cache(ttl, cached_json)
            log.debug("valid cache: %s, for url: %s", cache != [], url)
            endpoints.append((url, cache, cached_json, cached_zip, min_score, profile))

    # If we don't find our module in the yaml
    if not endpoints:
        return {}

    ret = {'Success': [], 'Failure': [], 'Controlled': []}
    # Dictionary of {pkg_name: list(pkg_versions)}
    local_pkgs = __salt__['pkg.list_pkgs'](versions_as_list=True)

    for url, cache, cached_json, cached_zip, min_score, profile in endpoints:
        log.debug("url: %s, min_score: %s", url, min_score)
        if cache: # Valid cached file
            master_json = cache
        else: # Query the url for cve's
            if url.startswith('http://') or url.startswith('https://'):
                if 'vulners.com' in url:
                    # Vulners api can only handles http:// requests from request.get
                    if url.startswith('https'):
                        url.replace('https', 'http', 1)
                    # Format the url for the request based on operating system.
                    if url.endswith('/'):
                        url = url[:-1]
                    url_final = '%s/api/v3/archive/distributive/?os=%s&version=%s' \
                                                                % (url, os_name, os_version)
                    log.debug('requesting: %s', url_final)
                    cve_query = requests.get(url_final)
                    # Confirm that the request was valid.
                    if cve_query.status_code != 200:
                        raise Exception('Vulners requests was not successful. Check the url.')
                    # Save vulners zip attachment in cache location and extract json
                    try:
                        with open(cached_zip, 'w') as zip_attachment:
                            zip_attachment.write(cve_query.content)
                        zip_file = ZipFile(cached_zip)
                        zip_file.extractall(os.path.dirname(cached_zip))
                        os.remove(cached_zip)
                        extracted_json = os.path.join(__opts__['cachedir'],
                                                      'cve_scan_cache',
                                                      '%s_%s.json' % (os_name, str(os_version).replace('.', '')))
                        log.debug('attempting to open %s', extracted_json)
                        with open(extracted_json, 'r') as json_file:
                            master_json = json.load(json_file)
                        os.remove(extracted_json)
                    except IOError as ioe:
                        log.error('The json zip attachment was not able to be extracted from vulners.')
                        raise ioe
                else: # Not a vulners request, external source for cve's
                    log.debug('requesting: %s', url)
                    cve_query = requests.get(url)
                    if cve_query.status_code != 200:
                        log.error('URL request was not successful.')
                        raise Exception('The url given is invalid.')
                    master_json = json.loads(cve_query.text)
                #Cache results.
                try:
                    with open(cached_json, 'w') as cache_file:
                        json.dump(master_json, cache_file)
                except IOError:
                    log.error('The cve results weren\'t able to be cached')
            elif url.startswith('salt://'):
                # Cache the file
                log.debug('getting file from %s', url)
                cache_file = __salt__['cp.get_file'](url, cached_json)
                if cache_file:
                    master_json = json.load(open(cache_file))
                else:
                    raise IOError('The file was not able to be retrieved from the salt file server.')
            else:
                raise Exception('The url is invalid. It does not begin with http(s):// or salt://')

        affected_pkgs = _get_cve_vulnerabilities(master_json, os_version)

        # Check all local packages against cve vulnerablities in affected_pkgs
        for local_pkg in local_pkgs:
            vulnerable = None
            if local_pkg in affected_pkgs:
                # There can be multiple versions for a single local package, check all
                for local_version in local_pkgs[local_pkg]:
                    # There can be multiple cve announcements for a single package, check against all
                    for affected_obj in affected_pkgs[local_pkg]:
                        affected_version = affected_obj.pkg_version
                        if _is_vulnerable(local_version, affected_version, affected_obj.operator):
                            # If the local pkg hasn't been found as vulnerable yet, vulnerable is None
                            if not vulnerable:
                                affected_obj.oudated_version = local_version
                                vulnerable = affected_obj
                            # If local_pkg has already been marked affected, vulnerable is set. We
                            #   want to report the cve with highest severity
                            else:
                                if affected_obj.score > vulnerable.score:
                                    affected_obj.oudated_version = local_version
                                    vulnerable = affected_obj
                if vulnerable:
                    if vulnerable.score < min_score:
                        ret['Controlled'].append(vulnerable.get_report(profile))
                    else:
                        ret['Failure'].append(vulnerable.get_report(profile))

    if tags != '*':
        log.debug("tags: %s", tags)
        remove = []
        for i, failure in enumerate(ret['Failure']):
            if not fnmatch.fnmatch(failure.keys()[0], tags):
                remove.append(i)
        remove.reverse()
        for i in remove:
            ret['Failure'].pop(i)

        remove = []
        for i, failure in enumerate(ret['Controlled']):
            if not fnmatch.fnmatch(failure.keys()[0], tags):
                remove.append(i)
        remove.reverse()
        for i in remove:
            ret['Controlled'].pop(i)

    if not ret['Controlled']:
        ret.pop('Controlled')

    return ret

def _get_cve_vulnerabilities(query_results, os_version):
    '''
    Returns dictionary of vulnerablities, mapped as pkg_name:pkgObj.
    '''

    vulnerable_pkgs = {}

    for report in query_results:
        try:
            reporter = report['_source'].get('reporter', '')
            cve_list = report['_source'].get('cvelist', [])
            href = report['_source'].get('href', '')
            score = report['_source']['cvss'].get('score', 0)
            title = report['_source'].get('title', 'No Title Given')

            for pkg in report['_source']['affectedPackage']:
                #_source:affectedPackages
                if pkg['OSVersion'] in ['any', os_version]: #Only use matching os
                    pkg_obj = VulnerablePkg(title, pkg['packageName'], pkg['packageVersion'], \
                                 score, pkg['operator'], reporter, href, cve_list)
                    if pkg_obj.pkg not in vulnerable_pkgs:
                        vulnerable_pkgs[pkg_obj.pkg] = [pkg_obj]
                    else:
                        vulnerable_pkgs[pkg_obj.pkg].append(pkg_obj)
        except KeyError, key_err:
            if key_err != '_source':
                log.error('Format error at: %s', report)
                raise KeyError('The cve data was not formatted correctly at: %s' % pkg)
            else:
                log.error('Format error at: %s', report)
                raise KeyError('The cve data was not formatted correctly')
    return vulnerable_pkgs

def _is_vulnerable(local_version, affected_version, operator):
    '''
    Given two version strings, and operator
        returns whether the package is vulnerable or not.
    '''
    # Get rid of prefix if version number has one, ex '1:3.4.52'
    if ':' in local_version:
        _, _, local_version = local_version.partition(':')
    if ':' in affected_version:
        _, _, affected_version = affected_version.partition(':')

    compare = None
    # Try to use salt's built in comparison module, if it exists for distro
    if 'pkg.version_cmp' in __salt__:
        compare = __salt__['pkg.version_cmp'](local_version, affected_version)

    # When salt can't compare, use LooseVersion
    if compare is None:
        #Compare from higher order to lower order based on '-' split.
        local_version_split = local_version.split('-')
        affected_version_split = affected_version.split('-')

        for (order_index, local_version_str) in enumerate(local_version_split):

            local_version_obj = LooseVersion(local_version_str)
            affected_version_obj = LooseVersion(affected_version_split[order_index])

            #Check lower order bits if higher order are equal.
            if local_version == affected_version:
                continue

            #Return when highest order version is not equal.
            elif local_version_obj > affected_version_obj:
                compare = 1
                break
            elif local_version_obj < affected_version_obj:
                compare = -1
                break
        # If for loop exits without break, the versions are equal.
        else:
            compare = 0

    # Return whether local_version is vulnerable to affected_verison
    if operator == 'le':
        return compare <= 0
    elif operator == 'lt':
        return compare < 0

def _get_cache(ttl, cache_path):
    '''
    If url contains valid cache, returns it, else returns empty list.
    '''
    # Check if we have a valid cached version.
    try:
        cached_time = os.path.getmtime(cache_path)
    except OSError:
        return []
    if current_time() - cached_time < ttl:
        log.debug('%s is less than ttl', cache_path)
        try:
            with open(cache_path) as json_file:
                loaded_json = json.load(json_file)
                return loaded_json
        except IOError:
            return []
        except ValueError:
            log.error('%s was not json formatted', cache_path)
            return []
    else:
        log.debug('%s was older than ttl', cache_path)
        return []

class VulnerablePkg:
    '''
    Object representing a vulnverable pkg for the current operating system.
    '''
    def __init__(self, title, pkg, pkg_version, score, operator, reporter, href, cve_list):
        self.title = title
        self.pkg = pkg
        self.pkg_version = pkg_version
        self.score = float(score)
        if operator not in ['lt', 'le']:
            log.error('pkg:%s contains an operator that\'s not supported and was changed to <')
            operator = 'lt'
        self.operator = operator
        self.href = href
        self.cve_list = cve_list
        self.reporter = reporter
        self.oudated_version = None

    def get_report(self, profile):
        '''
        Return the dictionary of what should be reported in failures, based on verbose.
        '''
        return {
            'tag': self.pkg + '-' + self.pkg_version,
            'href': self.href,
            'affected_version': self.pkg_version,
            'reporter': self.reporter,
            'score': self.score,
            'cve_list': self.cve_list,
            'affected_pkg': self.pkg,
            'local_version': self.oudated_version,
            'description': self.title,
            'nova_profile': profile
        }

<commit_msg>Fixes a mismatch of the os_version fetched from grains. Vulners.com report 'OSVersion' as found in 'osrelease' grain rather than as in 'osmajorrelease'.<commit_after># -*- encoding: utf-8 -*-
'''
HubbleStack Nova plugin for auditing installed packages.

This module checks all of a system's local packages and reports if the package
is vulnerable to a known cve. The cve vunlerablities are gathered via the url in
the yaml profile, and that data cached at the path
/var/cache/salt/minion/cve_scan_cache/<os_name>_<version>.json

:maintainer: HubbleStack / jaredhanson11
:maturity: 2016.7.0
:platform: Linux
:requires: SaltStack

This audit module requires yaml data to execute. It will search the local
directory for any .yaml files, and if it finds a top-level 'cve_scan_v2' key, it
will use that data.

Sample YAML data with inline comments:

cve_scan_v2:
    # Seconds until the local cache expires
    ttl: 86400
    # Source of cve data
    url: http://vulners.com/
    # Optional control tag
    control:
        # minimum score, vulnerabilities with a smaller
        #   score added to 'Controlled' output
        score: 3

The source of the cve data can be http://vulners.com/, salt://path/to/json, and
any other url that returns cve data in json format. If the url contains
vulners.com, then this module will use the local system's os and os version to
dynamically query vulner.com/api/v3 for cve data specifically related to your
system. If the url doesn't contain vulners.com, it will query the exact url, so
that endpoint must return cve data specific to the system you are scanning.

The cve data json must be formatted as follows:

[

{'_source': {'affectedPackage': [{'OS': 'CentOS',
                                    'OSVersion': '7',
                                    'operator': 'lt',
                                    'packageFilename': 'krb5-server-1.13.2-12.el7_2.x86_64.rpm',
                                    'packageName': 'krb5-server',
                                    'packageVersion': '1.13.2-12.el7_2'},
                                   {'OS': 'CentOS',
                                    'OSVersion': '7',
                                    'operator': 'lt',
                                    'packageFilename': 'krb5-libs-1.13.2-12.el7_2.i686.rpm',
                                    'packageName': 'krb5-libs',
                                    'packageVersion': '1.13.2-12.el7_2'}
                                    ]
              'cvelist': ['CVE-2015-8631',
                           'CVE-2015-8630',
                           'CVE-2015-8629'],
              'cvss': {'score': 6.8}
              'href': 'http://lists.centos.org/pipermail/centos-announce/2016-March/021788.html',
              'reporter': 'CentOS Project',
              'title': 'Moderate krb5 Security Update'
            }
    },
...

 ]

'''
from __future__ import absolute_import
import logging

import fnmatch
import hashlib
import json
import os
import re
import requests

from distutils.version import LooseVersion
from time import time as current_time
from zipfile import ZipFile

import salt
import salt.utils

log = logging.getLogger(__name__)

def __virtual__():
    return not salt.utils.is_windows()

def audit(data_list, tags, debug=False, **kwargs):
    '''
    Main audit function. See module docstring for more information on usage.
    '''
    os_version = __grains__.get('osrelease', None)
    if os_version is None:
        os_version = __grains__.get('osmajorrelease', None)
    os_name = __grains__['os'].lower()

    log.debug("os_version: %s, os_name: %s", os_version, os_name)

    endpoints = []

    # Go through yaml to check for cve_scan_v2,
    #    if its present, check for a cached version
    #    of the scan.
    for profile, data in data_list:

        if 'cve_scan_v2' in data:

            ttl = data['cve_scan_v2']['ttl']
            url = data['cve_scan_v2']['url']
            control = data['cve_scan_v2'].get('control', {})
            # Ability to add more controls easily, in control dict
            min_score = float(control.get('score', 0))
            urlhash = hashlib.md5(url).hexdigest()
            cached_json = os.path.join(__opts__['cachedir'],
                                       'cve_scan_cache',
                                       '%s.json' % urlhash)
            cached_zip = os.path.join(__opts__['cachedir'],
                                      'cve_scan_cache',
                                      '%s.zip' % urlhash)
            # Make cache directory and all parent directories if it doesn't exist.
            if not os.path.exists(os.path.dirname(cached_json)):
                os.makedirs(os.path.dirname(cached_json))
            cache = _get_cache(ttl, cached_json)
            log.debug("valid cache: %s, for url: %s", cache != [], url)
            endpoints.append((url, cache, cached_json, cached_zip, min_score, profile))

    # If we don't find our module in the yaml
    if not endpoints:
        return {}

    ret = {'Success': [], 'Failure': [], 'Controlled': []}
    # Dictionary of {pkg_name: list(pkg_versions)}
    local_pkgs = __salt__['pkg.list_pkgs'](versions_as_list=True)

    for url, cache, cached_json, cached_zip, min_score, profile in endpoints:
        log.debug("url: %s, min_score: %s", url, min_score)
        if cache: # Valid cached file
            master_json = cache
        else: # Query the url for cve's
            if url.startswith('http://') or url.startswith('https://'):
                if 'vulners.com' in url:
                    # Vulners api can only handles http:// requests from request.get
                    if url.startswith('https'):
                        url.replace('https', 'http', 1)
                    # Format the url for the request based on operating system.
                    if url.endswith('/'):
                        url = url[:-1]
                    url_final = '%s/api/v3/archive/distributive/?os=%s&version=%s' \
                                                                % (url, os_name, os_version)
                    log.debug('requesting: %s', url_final)
                    cve_query = requests.get(url_final)
                    # Confirm that the request was valid.
                    if cve_query.status_code != 200:
                        raise Exception('Vulners requests was not successful. Check the url.')
                    # Save vulners zip attachment in cache location and extract json
                    try:
                        with open(cached_zip, 'w') as zip_attachment:
                            zip_attachment.write(cve_query.content)
                        zip_file = ZipFile(cached_zip)
                        zip_file.extractall(os.path.dirname(cached_zip))
                        os.remove(cached_zip)
                        extracted_json = os.path.join(__opts__['cachedir'],
                                                      'cve_scan_cache',
                                                      '%s_%s.json' % (os_name, str(os_version).replace('.', '')))
                        log.debug('attempting to open %s', extracted_json)
                        with open(extracted_json, 'r') as json_file:
                            master_json = json.load(json_file)
                        os.remove(extracted_json)
                    except IOError as ioe:
                        log.error('The json zip attachment was not able to be extracted from vulners.')
                        raise ioe
                else: # Not a vulners request, external source for cve's
                    log.debug('requesting: %s', url)
                    cve_query = requests.get(url)
                    if cve_query.status_code != 200:
                        log.error('URL request was not successful.')
                        raise Exception('The url given is invalid.')
                    master_json = json.loads(cve_query.text)
                #Cache results.
                try:
                    with open(cached_json, 'w') as cache_file:
                        json.dump(master_json, cache_file)
                except IOError:
                    log.error('The cve results weren\'t able to be cached')
            elif url.startswith('salt://'):
                # Cache the file
                log.debug('getting file from %s', url)
                cache_file = __salt__['cp.get_file'](url, cached_json)
                if cache_file:
                    master_json = json.load(open(cache_file))
                else:
                    raise IOError('The file was not able to be retrieved from the salt file server.')
            else:
                raise Exception('The url is invalid. It does not begin with http(s):// or salt://')

        affected_pkgs = _get_cve_vulnerabilities(master_json, os_version)

        # Check all local packages against cve vulnerablities in affected_pkgs
        for local_pkg in local_pkgs:
            vulnerable = None
            if local_pkg in affected_pkgs:
                # There can be multiple versions for a single local package, check all
                for local_version in local_pkgs[local_pkg]:
                    # There can be multiple cve announcements for a single package, check against all
                    for affected_obj in affected_pkgs[local_pkg]:
                        affected_version = affected_obj.pkg_version
                        if _is_vulnerable(local_version, affected_version, affected_obj.operator):
                            # If the local pkg hasn't been found as vulnerable yet, vulnerable is None
                            if not vulnerable:
                                affected_obj.oudated_version = local_version
                                vulnerable = affected_obj
                            # If local_pkg has already been marked affected, vulnerable is set. We
                            #   want to report the cve with highest severity
                            else:
                                if affected_obj.score > vulnerable.score:
                                    affected_obj.oudated_version = local_version
                                    vulnerable = affected_obj
                if vulnerable:
                    if vulnerable.score < min_score:
                        ret['Controlled'].append(vulnerable.get_report(profile))
                    else:
                        ret['Failure'].append(vulnerable.get_report(profile))

    if tags != '*':
        log.debug("tags: %s", tags)
        remove = []
        for i, failure in enumerate(ret['Failure']):
            if not fnmatch.fnmatch(failure.keys()[0], tags):
                remove.append(i)
        remove.reverse()
        for i in remove:
            ret['Failure'].pop(i)

        remove = []
        for i, failure in enumerate(ret['Controlled']):
            if not fnmatch.fnmatch(failure.keys()[0], tags):
                remove.append(i)
        remove.reverse()
        for i in remove:
            ret['Controlled'].pop(i)

    if not ret['Controlled']:
        ret.pop('Controlled')

    return ret

def _get_cve_vulnerabilities(query_results, os_version):
    '''
    Returns dictionary of vulnerablities, mapped as pkg_name:pkgObj.
    '''

    vulnerable_pkgs = {}

    for report in query_results:
        try:
            reporter = report['_source'].get('reporter', '')
            cve_list = report['_source'].get('cvelist', [])
            href = report['_source'].get('href', '')
            score = report['_source']['cvss'].get('score', 0)
            title = report['_source'].get('title', 'No Title Given')

            for pkg in report['_source']['affectedPackage']:
                #_source:affectedPackages
                if pkg['OSVersion'] in ['any', os_version]: #Only use matching os
                    pkg_obj = VulnerablePkg(title, pkg['packageName'], pkg['packageVersion'], \
                                 score, pkg['operator'], reporter, href, cve_list)
                    if pkg_obj.pkg not in vulnerable_pkgs:
                        vulnerable_pkgs[pkg_obj.pkg] = [pkg_obj]
                    else:
                        vulnerable_pkgs[pkg_obj.pkg].append(pkg_obj)
        except KeyError, key_err:
            if key_err != '_source':
                log.error('Format error at: %s', report)
                raise KeyError('The cve data was not formatted correctly at: %s' % pkg)
            else:
                log.error('Format error at: %s', report)
                raise KeyError('The cve data was not formatted correctly')
    return vulnerable_pkgs

def _is_vulnerable(local_version, affected_version, operator):
    '''
    Given two version strings, and operator
        returns whether the package is vulnerable or not.
    '''
    # Get rid of prefix if version number has one, ex '1:3.4.52'
    if ':' in local_version:
        _, _, local_version = local_version.partition(':')
    if ':' in affected_version:
        _, _, affected_version = affected_version.partition(':')

    compare = None
    # Try to use salt's built in comparison module, if it exists for distro
    if 'pkg.version_cmp' in __salt__:
        compare = __salt__['pkg.version_cmp'](local_version, affected_version)

    # When salt can't compare, use LooseVersion
    if compare is None:
        #Compare from higher order to lower order based on '-' split.
        local_version_split = local_version.split('-')
        affected_version_split = affected_version.split('-')

        for (order_index, local_version_str) in enumerate(local_version_split):

            local_version_obj = LooseVersion(local_version_str)
            affected_version_obj = LooseVersion(affected_version_split[order_index])

            #Check lower order bits if higher order are equal.
            if local_version == affected_version:
                continue

            #Return when highest order version is not equal.
            elif local_version_obj > affected_version_obj:
                compare = 1
                break
            elif local_version_obj < affected_version_obj:
                compare = -1
                break
        # If for loop exits without break, the versions are equal.
        else:
            compare = 0

    # Return whether local_version is vulnerable to affected_verison
    if operator == 'le':
        return compare <= 0
    elif operator == 'lt':
        return compare < 0

def _get_cache(ttl, cache_path):
    '''
    If url contains valid cache, returns it, else returns empty list.
    '''
    # Check if we have a valid cached version.
    try:
        cached_time = os.path.getmtime(cache_path)
    except OSError:
        return []
    if current_time() - cached_time < ttl:
        log.debug('%s is less than ttl', cache_path)
        try:
            with open(cache_path) as json_file:
                loaded_json = json.load(json_file)
                return loaded_json
        except IOError:
            return []
        except ValueError:
            log.error('%s was not json formatted', cache_path)
            return []
    else:
        log.debug('%s was older than ttl', cache_path)
        return []

class VulnerablePkg:
    '''
    Object representing a vulnverable pkg for the current operating system.
    '''
    def __init__(self, title, pkg, pkg_version, score, operator, reporter, href, cve_list):
        self.title = title
        self.pkg = pkg
        self.pkg_version = pkg_version
        self.score = float(score)
        if operator not in ['lt', 'le']:
            log.error('pkg:%s contains an operator that\'s not supported and was changed to <')
            operator = 'lt'
        self.operator = operator
        self.href = href
        self.cve_list = cve_list
        self.reporter = reporter
        self.oudated_version = None

    def get_report(self, profile):
        '''
        Return the dictionary of what should be reported in failures, based on verbose.
        '''
        return {
            'tag': self.pkg + '-' + self.pkg_version,
            'href': self.href,
            'affected_version': self.pkg_version,
            'reporter': self.reporter,
            'score': self.score,
            'cve_list': self.cve_list,
            'affected_pkg': self.pkg,
            'local_version': self.oudated_version,
            'description': self.title,
            'nova_profile': profile
        }

<|endoftext|>


<commit_before># -*- encoding: utf-8 -*-
'''
Hubble Nova plugin for FreeBSD pkgng audit

:maintainer: HubbleStack
:maturity: 20160421
:platform: FreeBSD
:requires: SaltStack
'''
from __future__ import absolute_import
import logging

log = logging.getLogger(__name__)

__tags__ = None

def __virtual__():
    if 'FreeBSD' not in __grains__['os']:
        return False, 'This audit module only runs on FreeBSD'
    global __tags__
    __tags__ = ['freebsd-pkg-audit']
    return True

def audit(tags, verbose=False):
    '''
    Run the pkg.audit command
    '''
    ret = {'Success': [], 'Failure': []}

    salt_ret = __salt__['pkg.audit']()
    if '0 problem(s)' not in salt_ret:
        ret['Failure'].append(salt_ret)
    else:
        ret['Success'].append(salt_ret)

    return ret
<commit_msg>Update frebsd-pkg-audit to rely on yaml data and take data from hubble.py<commit_after># -*- encoding: utf-8 -*-
'''
Hubble Nova plugin for FreeBSD pkgng audit

:maintainer: HubbleStack
:maturity: 20160421
:platform: FreeBSD
:requires: SaltStack
'''
from __future__ import absolute_import
import logging

log = logging.getLogger(__name__)

def __virtual__():
    if 'FreeBSD' not in __grains__['os']:
        return False, 'This audit module only runs on FreeBSD'
    return True

def audit(data_list, tags, verbose=False):
    '''
    Run the pkg.audit command
    '''
    ret = {'Success': [], 'Failure': []}

    __tags__ = []
    for data in data_list:
        if 'freebsd-pkg' in data:
            __tags__ = ['freebsd-pkg-audit']
            break

    if not __tags__:
        # No yaml data found, don't do any work
        return ret

    salt_ret = __salt__['pkg.audit']()
    if '0 problem(s)' not in salt_ret:
        ret['Failure'].append(salt_ret)
    else:
        ret['Success'].append(salt_ret)

    return ret
<|endoftext|>


<commit_before># The MIT License (MIT) # Copyright (c) 2014-2017 University of Bristol
#
#  Permission is hereby granted, free of charge, to any person obtaining a copy
#  of this software and associated documentation files (the "Software"), to deal
#  in the Software without restriction, including without limitation the rights
#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#  copies of the Software, and to permit persons to whom the Software is
#  furnished to do so, subject to the following conditions:
#
#  The above copyright notice and this permission notice shall be included in all
#  copies or substantial portions of the Software.
#
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
#  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
#  MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
#  IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
#  DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
#  OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE
#  OR OTHER DEALINGS IN THE SOFTWARE.

import pytz
from datetime import datetime, timedelta

class UTC(pytz.UTC):
    def __repr__(self):
        return "UTC"

MIN_DATE = datetime.min.replace(tzinfo=UTC)
MAX_DATE = datetime.max.replace(tzinfo=UTC).replace(microsecond=0)

def utcnow():
    return datetime.utcnow().replace(tzinfo=UTC)

def get_timedelta(value):
    if isinstance(value, (int, float)):
        return timedelta(seconds=value)
    elif isinstance(value, timedelta):
        return value
    else:
        raise ValueError("Expected int, float, or timedelta, got {}".format(type(value)))

def unix2datetime(u):
    return datetime.fromtimestamp(u / 1000.0, tz=UTC) + timedelta(hours=0)

def duration2str(x):
    minutes, seconds = divmod(x.total_seconds(), 60)
    return '{} min {} sec'.format(int(minutes), int(seconds))

def construct_experiment_id(time_interval):
    """
    Construct an experiment id from a time interval
    :param instance:
    :return: The experiment id
    :type time_interval: TimeInterval
    :rtype: str
    """
    # Construct id based on unix epoch timestamps
    epoch = datetime.utcfromtimestamp(0).replace(tzinfo=UTC)
    start = int((time_interval.start - epoch).total_seconds() * 1000.0)
    end = int((time_interval.end - epoch).total_seconds() * 1000.0)
    return "{}-{}".format(start, end)
<commit_msg>Updated utcnow to be only millisecond precise<commit_after># The MIT License (MIT) # Copyright (c) 2014-2017 University of Bristol
#
#  Permission is hereby granted, free of charge, to any person obtaining a copy
#  of this software and associated documentation files (the "Software"), to deal
#  in the Software without restriction, including without limitation the rights
#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#  copies of the Software, and to permit persons to whom the Software is
#  furnished to do so, subject to the following conditions:
#
#  The above copyright notice and this permission notice shall be included in all
#  copies or substantial portions of the Software.
#
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
#  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
#  MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
#  IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
#  DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
#  OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE
#  OR OTHER DEALINGS IN THE SOFTWARE.

import pytz
from datetime import datetime, timedelta

class UTC(pytz.UTC):
    def __repr__(self):
        return "UTC"

MIN_DATE = datetime.min.replace(tzinfo=UTC)
MAX_DATE = datetime.max.replace(tzinfo=UTC).replace(microsecond=0)

def utcnow():
    """
    Gets the current datetime in UTC format with millisecond precision
    :return:
    """
    now = datetime.utcnow().replace(tzinfo=UTC)
    return datetime(now.year, now.month, now.day, now.hour, now.minute, now.second, now.microsecond / 1000 * 1000, UTC)

def get_timedelta(value):
    if isinstance(value, (int, float)):
        return timedelta(seconds=value)
    elif isinstance(value, timedelta):
        return value
    else:
        raise ValueError("Expected int, float, or timedelta, got {}".format(type(value)))

def unix2datetime(u):
    return datetime.fromtimestamp(u / 1000.0, tz=UTC) + timedelta(hours=0)

def duration2str(x):
    minutes, seconds = divmod(x.total_seconds(), 60)
    return '{} min {} sec'.format(int(minutes), int(seconds))

def construct_experiment_id(time_interval):
    """
    Construct an experiment id from a time interval
    :param instance:
    :return: The experiment id
    :type time_interval: TimeInterval
    :rtype: str
    """
    # Construct id based on unix epoch timestamps
    epoch = datetime.utcfromtimestamp(0).replace(tzinfo=UTC)
    start = int((time_interval.start - epoch).total_seconds() * 1000.0)
    end = int((time_interval.end - epoch).total_seconds() * 1000.0)
    return "{}-{}".format(start, end)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: latin-1 -*-
#
#   Copyright 2016-2019 Blaise Frederick
#
#   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.
#
#
# $Author: frederic $
# $Date: 2016/07/11 14:50:43 $
# $Id: rapidtide,v 1.161 2016/07/11 14:50:43 frederic Exp $
#
#
#

from __future__ import print_function, division

import gc

import numpy as np

import rapidtide.multiproc as tide_multiproc
import rapidtide.resample as tide_resample
import rapidtide.util as tide_util

# this is here until numpy deals with their fft issue
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)

def _procOneVoxelPeaks(vox,
                       thetc,
                       themutualinformationator,
                       fmri_x,
                       fmritc,
                       os_fmri_x,
                       xcorr_x,
                       thexcorr,
                       lagmin,
                       lagmax,
                       oversampfactor=1,
                       interptype='univariate'
                       ):

    if oversampfactor >= 1:
        thetc[:] = tide_resample.doresample(fmri_x, fmritc, os_fmri_x, method=interptype)
    else:
        thetc[:] = fmritc
    thepeaks = tide_fit.getpeaks(xcorr_x, thexcorr, xrange=(lagmin, lagmax), display=True)
    peaklocs = []
    for thepeak in thepeaks:
        peaklocs.append(int(round(thepeak[2], 0)))
    theMI_list = themutualinformationator.run(thetc, locs=peaklocs)
    hybridpeaks = []
    for i in range(len(thepeaks)):
        hybridpeaks.append([thepeaks[i][0], thepeaks[i][1], theMI_list[i]])
    return vox, hybridpeaks

def peakevalpass(
        fmridata,
        referencetc,
        themutualinformationator,
        fmri_x,
        os_fmri_x,
        lagmininpts,
        lagmaxinpts,
        corrout,
        meanval,
        nprocs=1,
        oversampfactor=1,
        interptype='univariate',
        showprogressbar=True,
        chunksize=1000):
    """

    Parameters
    ----------
    fmridata
    referencetc
    thecorrelator
    fmri_x
    os_fmri_x
    tr
    lagmininpts
    lagmaxinpts
    corrout
    meanval
    nprocs
    oversampfactor
    interptype
    showprogressbar
    chunksize
    rt_floatset
    rt_floattype

    Returns
    -------

    """
    themutualinformationator.setreftc(referencetc)
    themutualinformationator.setlimits(lagmininpts, lagmaxinpts)

    inputshape = np.shape(fmridata)
    volumetotal = 0
    reportstep = 1000
    thetc = np.zeros(np.shape(os_fmri_x), dtype=rt_floattype)
    theglobalmaxlist = []
    if nprocs > 1:
        # define the consumer function here so it inherits most of the arguments
        def correlation_consumer(inQ, outQ):
            while True:
                try:
                    # get a new message
                    val = inQ.get()

                    # this is the 'TERM' signal
                    if val is None:
                        break

                    # process and send the data
                    outQ.put(_procOneVoxelPeaks(
                        val,
                        thetc,
                        themutualinformationator,
                        fmri_x,
                        fmridata[val, :],
                        os_fmri_x,
                        xcorr_x,
                        thexcorr,
                        lagmin,
                        lagmax,
                        oversampfactor=oversampfactor,
                        interptype=interptype)
                    )

                except Exception as e:
                    print("error!", e)
                    break

        data_out = tide_multiproc.run_multiproc(
            correlation_consumer,
            inputshape, None,
            nprocs=nprocs,
            showprogressbar=showprogressbar,
            chunksize=chunksize)

        # unpack the data
        volumetotal = 0
        for voxel in data_out:
            meanval[voxel[0]] = voxel[1]
            corrout[voxel[0], :] = voxel[2]
            thecorrscale = voxel[3]
            theglobalmaxlist.append(voxel[4] + 0)
            volumetotal += 1
        del data_out
    else:
        for vox in range(0, inputshape[0]):
            if (vox % reportstep == 0 or vox == inputshape[0] - 1) and showprogressbar:
                tide_util.progressbar(vox + 1, inputshape[0], label='Percent complete')
            dummy, thepeaks = \
                _procOneVoxelPeaks(
                    vox,
                    thetc,
                    themutualinformationator,
                    fmri_x,
                    fmridata[val, :],
                    os_fmri_x,
                    xcorr_x,
                    thexcorr,
                    lagmin,
                    lagmax,
                    oversampfactor=oversampfactor,
                    interptype=interptype)
            theglobalmaxlist.append(theglobalmax + 0)
            volumetotal += 1
    print('\nPeak evaluation performed on ' + str(volumetotal) + ' voxels')

    # garbage collect
    collected = gc.collect()
    print("Garbage collector: collected %d objects." % collected)

    return volumetotal, theglobalmaxlist, thecorrscale
<commit_msg>Process full range of correlation (limit on input)<commit_after>#!/usr/bin/env python
# -*- coding: latin-1 -*-
#
#   Copyright 2016-2019 Blaise Frederick
#
#   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.
#
#
# $Author: frederic $
# $Date: 2016/07/11 14:50:43 $
# $Id: rapidtide,v 1.161 2016/07/11 14:50:43 frederic Exp $
#
#
#

from __future__ import print_function, division

import gc

import numpy as np

import rapidtide.multiproc as tide_multiproc
import rapidtide.resample as tide_resample
import rapidtide.util as tide_util
import rapidtide.fit as tide_fit

# this is here until numpy deals with their fft issue
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)

def _procOneVoxelPeaks(vox,
                       thetc,
                       themutualinformationator,
                       fmri_x,
                       fmritc,
                       os_fmri_x,
                       xcorr_x,
                       thexcorr,
                       oversampfactor=1,
                       sort=True,
                       interptype='univariate'
                       ):

    if oversampfactor >= 1:
        thetc[:] = tide_resample.doresample(fmri_x, fmritc, os_fmri_x, method=interptype)
    else:
        thetc[:] = fmritc
    thepeaks = tide_fit.getpeaks(xcorr_x, thexcorr, display=False)
    peaklocs = []
    for thepeak in thepeaks:
        peaklocs.append(int(round(thepeak[2], 0)))
    theMI_list = themutualinformationator.run(thetc, locs=peaklocs)
    hybridpeaks = []
    for i in range(len(thepeaks)):
        hybridpeaks.append([thepeaks[i][0], thepeaks[i][1], theMI_list[i]])
    if sort:
        hybridpeaks.sort(key=lambda x: x[2], reverse=True)
    return vox, hybridpeaks

def peakevalpass(
        fmridata,
        corrdata,
        referencetc,
        themutualinformationator,
        fmri_x,
        os_fmri_x,
        nprocs=1,
        oversampfactor=1,
        interptype='univariate',
        showprogressbar=True,
        chunksize=1000):
    """

    Parameters
    ----------
    fmridata
    referencetc
    thecorrelator
    fmri_x
    os_fmri_x
    tr
    corrout
    meanval
    nprocs
    oversampfactor
    interptype
    showprogressbar
    chunksize
    rt_floatset
    rt_floattype

    Returns
    -------

    """
    peakdict = {}
    themutualinformationator.setreftc(referencetc)
    #themutualinformationator.setlimits(lagmininpts, lagmaxinpts)

    inputshape = np.shape(fmridata)
    volumetotal = 0
    reportstep = 1000
    thetc = np.zeros(np.shape(os_fmri_x), dtype=rt_floattype)
    theglobalmaxlist = []
    if nprocs > 1:
        # define the consumer function here so it inherits most of the arguments
        def correlation_consumer(inQ, outQ):
            while True:
                try:
                    # get a new message
                    val = inQ.get()

                    # this is the 'TERM' signal
                    if val is None:
                        break

                    # process and send the data
                    outQ.put(_procOneVoxelPeaks(
                        val,
                        thetc,
                        themutualinformationator,
                        fmri_x,
                        fmridata[val, :],
                        os_fmri_x,
                        xcorr_x,
                        corrdata[val, :],
                        oversampfactor=oversampfactor,
                        interptype=interptype)
                    )

                except Exception as e:
                    print("error!", e)
                    break

        data_out = tide_multiproc.run_multiproc(
            correlation_consumer,
            inputshape, None,
            nprocs=nprocs,
            showprogressbar=showprogressbar,
            chunksize=chunksize)

        # unpack the data
        volumetotal = 0
        for voxel in data_out:
            peakdict[str(voxel[0])] = voxel[1]
            volumetotal += 1
        del data_out
    else:
        for vox in range(0, inputshape[0]):
            if (vox % reportstep == 0 or vox == inputshape[0] - 1) and showprogressbar:
                tide_util.progressbar(vox + 1, inputshape[0], label='Percent complete')
            dummy, peakdict[str(vox)] = \
                _procOneVoxelPeaks(
                    vox,
                    thetc,
                    themutualinformationator,
                    fmri_x,
                    fmridata[val, :],
                    os_fmri_x,
                    xcorr_x,
                    thexcorr,
                    oversampfactor=oversampfactor,
                    interptype=interptype)
            theglobalmaxlist.append(theglobalmax + 0)
            volumetotal += 1
    print('\nPeak evaluation performed on ' + str(volumetotal) + ' voxels')

    # garbage collect
    collected = gc.collect()
    print("Garbage collector: collected %d objects." % collected)

    return volumetotal, theglobalmaxlist, thecorrscale
<|endoftext|>


<commit_before>from rdflib.Store import Store
from rdflib.Store import check_subject, check_predicate, check_object

from rdflib.backends.InMemoryBackend import InMemoryBackend
from rdflib.backends.Concurrent import Concurrent

class TripleStore(Store):

    def __init__(self, location=None, backend=None):
        if backend==None:
            #backend = Concurrent(InMemoryBackend())
            backend = InMemoryBackend()
        super(TripleStore, self).__init__(backend)
        if location:
            self.load(location)

    def add(self, (subject, predicate, object)):
        check_subject(subject)
        check_predicate(predicate)
        check_object(object)
        self.backend.add((subject, predicate, object))

    def remove(self, (subject, predicate, object)):
        if subject:
            check_subject(subject)
        if predicate:
            check_predicate(predicate)
        if object:
            check_object(object)
        self.backend.remove((subject, predicate, object))

    def triples(self, (subject, predicate, object)):
        if subject:
            check_subject(subject)
        if predicate:
            check_predicate(predicate)
        if object:
            check_object(object)
        for triple in self.backend.triples((subject, predicate, object)):
            yield triple

<commit_msg>added missing from __future__<commit_after>from __future__ import generators
from rdflib.Store import Store
from rdflib.Store import check_subject, check_predicate, check_object

from rdflib.backends.InMemoryBackend import InMemoryBackend
from rdflib.backends.Concurrent import Concurrent

class TripleStore(Store):

    def __init__(self, location=None, backend=None):
        if backend==None:
            #backend = Concurrent(InMemoryBackend())
            backend = InMemoryBackend()
        super(TripleStore, self).__init__(backend)
        if location:
            self.load(location)

    def add(self, (subject, predicate, object)):
        check_subject(subject)
        check_predicate(predicate)
        check_object(object)
        self.backend.add((subject, predicate, object))

    def remove(self, (subject, predicate, object)):
        if subject:
            check_subject(subject)
        if predicate:
            check_predicate(predicate)
        if object:
            check_object(object)
        self.backend.remove((subject, predicate, object))

    def triples(self, (subject, predicate, object)):
        if subject:
            check_subject(subject)
        if predicate:
            check_predicate(predicate)
        if object:
            check_object(object)
        for triple in self.backend.triples((subject, predicate, object)):
            yield triple

<|endoftext|>