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

import os.path as osp

from pibooth import fonts
from pibooth import pictures

from pibooth.language import get_translated_text

ARROW_TOP = 'top'
ARROW_BOTTOM = 'bottom'
ARROW_HIDDEN = 'hidden'

class Background(object):

    def __init__(self, image_name, color=(0, 0, 0), text_color=(255, 255, 255)):
        self._rect = None
        self._need_update = False

        self._background = None
        self._background_color = color
        self._background_image = None

        self._overlay = None
        self._overlay_image = "{0}.png".format(image_name)
        self.text = get_translated_text(image_name)
        self.text_size = 72
        self.text_x = 0
        self.text_y = 0
        self._text_color = text_color

    def __str__(self):
        """Return background final name.
        """
        return self._overlay_image

    def get_rect(self):
        """Return the Rect object of the background image. As aspect ratio
        is kept, the size of the background may be different from the
        screen size.
        """
        return self._overlay.get_rect(center=self._rect.center)

    def set_color(self, color_or_path):
        """Set background color (RGB tuple) or path to an image that used to
        fill the background.

        :param color_or_path: RGB color tuple or image path
        :type color_or_path: tuple or str
        """
        if isinstance(color_or_path, (tuple, list)):
            assert len(color_or_path) == 3, "Length of 3 is required for RGB tuple"
            if color_or_path != self._background_color:
                self._background_color = color_or_path
                self._need_update = True
        else:
            assert osp.isfile(color_or_path), "Invalid image for window background: '{}'".format(color_or_path)
            if color_or_path != self._background_image:
                self._background_image = color_or_path
                self._need_update = True

    def set_text_color(self, color):
        """Set text color (RGB tuple) used to write the texts.
        :param color: RGB color tuple
        :type color: tuple
        """
        assert len(color) == 3, "Length of 3 is required for RGB tuple"
        if color != self._text_color:
            self._text_color = color
            self._need_update = True

    def resize(self, screen):
        """Resize objects to fit to the screen.
        """
        if self._rect != screen.get_rect():
            self._rect = screen.get_rect()
            self.text_size = 72*self._rect.height//400
            self._overlay = pictures.get_pygame_image(self._overlay_image, (self._rect.width, self._rect.height))
            if self._background_image:
                self._background = pictures.get_pygame_image(
                    self._background_image, (self._rect.width, self._rect.height), crop=True)
            self._need_update = True

    def paint(self, screen):
        """Paint and animate the surfaces on the screen.
        """
        if self._background:
            screen.blit(self._background, (0, 0))
        else:
            screen.fill(self._background_color)
        screen.blit(self._overlay, self.get_rect())
        self.write_text(screen)
        self._need_update = False

    def write_text(self, screen):
        """Add the text to the screen
        """
        text_font = pictures.pygame.font.Font(fonts.get_filename("Amatic-Bold"), self.text_size)
        text = text_font.render(self.text, True, self._text_color, (0, 0, 0))
        textrect = text.get_rect()
        textrect.centerx = self.text_x
        textrect.centery = self.text_y
        black = (0, 0, 0)
        text.set_colorkey(black)  # Black colors will not be blit.
        screen.blit(text, textrect)
        self._need_update = False

class IntroBackground(Background):

    def __init__(self, arrow_location=ARROW_BOTTOM, arrow_offset=0):
        Background.__init__(self, "intro")
        self.arrow_location = arrow_location
        self.arrow_offset = arrow_offset
        self.left_arrow = None
        self.left_arrow_pos = None

    def resize(self, screen):
        Background.resize(self, screen)
        self.text_x = screen.get_rect().centerx/2
        self.text_y = screen.get_rect().centery
        if self._need_update:
            if self.arrow_location != ARROW_HIDDEN:
                size = (self.get_rect().width * 0.3, self.get_rect().height * 0.3)

                vflip = True if self.arrow_location == ARROW_TOP else False
                self.left_arrow = pictures.get_pygame_image("arrow.png", size, vflip=vflip)

                x = int(self.get_rect().left + self.get_rect().width // 4
                        - self.left_arrow.get_rect().width // 2)
                if self.arrow_location == ARROW_TOP:
                    y = self.get_rect().top + 10
                else:
                    y = int(self.get_rect().top + 2 * self.get_rect().height // 3)

                self.left_arrow_pos = (x - self.arrow_offset, y)

    def paint(self, screen):
        Background.paint(self, screen)
        if self.arrow_location != ARROW_HIDDEN:
            screen.blit(self.left_arrow, self.left_arrow_pos)

class IntroWithPrintBackground(IntroBackground):

    def __init__(self, arrow_location=ARROW_BOTTOM, arrow_offset=0):
        IntroBackground.__init__(self, arrow_location, arrow_offset)
        self._overlay_image = "intro_print"
        self.right_arrow = None
        self.right_arrow_pos = None

    def write_text(self, screen):
        IntroBackground.write_text(self, screen)
        text_size = 20*screen.get_rect().height//400
        text_font = pictures.pygame.font.Font(fonts.get_filename("Amatic-Bold"), text_size)
        text_strings = get_translated_text("intro_print").splitlines()
        delta_y = 0
        for text_string in text_strings:
            text = text_font.render(text_string, True, (255, 255, 255), (0, 0, 0))
            textrect = text.get_rect()
            textrect.centerx = self.get_rect().width * 45/100
            textrect.centery = self.get_rect().height * 2/3 + delta_y
            black = (0, 0, 0)
            text.set_colorkey(black)  # Black colors will not be blit.
            screen.blit(text, textrect)
            delta_y += text.get_height()
        self._need_update = False

    def resize(self, screen):
        IntroBackground.resize(self, screen)
        if self.arrow_location != ARROW_HIDDEN:
            size = (self.get_rect().width * 0.1, self.get_rect().height * 0.1)
            vflip = True if self.arrow_location == ARROW_TOP else False
            self.right_arrow = pictures.get_pygame_image("arrow.png", size, hflip=True, vflip=vflip)

            x = int(self.get_rect().left + self.get_rect().width // 2
                    - self.right_arrow.get_rect().width // 2)
            if self.arrow_location == ARROW_TOP:
                y = self.get_rect().top + 10
            else:
                y = int(self.get_rect().top + 8 * self.get_rect().height // 9)
            self.right_arrow_pos = (x - self.arrow_offset, y)

    def paint(self, screen):
        IntroBackground.paint(self, screen)
        if self.arrow_location != ARROW_HIDDEN:
            screen.blit(self.right_arrow, self.right_arrow_pos)

class ChooseBackground(Background):

    def __init__(self, choices, arrow_location=ARROW_BOTTOM, arrow_offset=0):
        Background.__init__(self, "choose")
        self.arrow_location = arrow_location
        self.arrow_offset = arrow_offset
        self.choices = choices
        self.layout0 = None
        self.layout0_pos = None
        self.layout1 = None
        self.layout1_pos = None
        self.left_arrow = None
        self.left_arrow_pos = None
        self.right_arrow = None
        self.right_arrow_pos = None

    def resize(self, screen):
        Background.resize(self, screen)
        self.text_x = screen.get_rect().centerx
        self.text_y = screen.get_rect().height/7
        if self._need_update:
            size = (self.get_rect().width * 0.6, self.get_rect().height * 0.6)
            self.layout0 = pictures.get_pygame_image("layout{}.png".format(self.choices[0]), size)
            self.layout1 = pictures.get_pygame_image("layout{}.png".format(self.choices[1]), size)

            inter = (self.get_rect().width - 2 * self.layout0.get_rect().width) // 3

            x0 = int(self.get_rect().left + inter)
            x1 = int(self.get_rect().left + 2 * inter + self.layout0.get_rect().width)
            y = int(self.get_rect().top + self.get_rect().height * 0.3)

            self.layout0_pos = (x0, y)
            self.layout1_pos = (x1, y)

            if self.arrow_location != ARROW_HIDDEN:
                if self.arrow_location == ARROW_TOP:
                    y = self.get_rect().top + 5
                    x_offset = 30
                    size = (self.get_rect().width * 0.1, self.layout0_pos[1] - y - 10)
                else:
                    x_offset = 0
                    y = self.layout0_pos[1] + self.layout0.get_rect().height + 5
                    size = (self.get_rect().width * 0.1, self.get_rect().bottom - y - 5)

                vflip = True if self.arrow_location == ARROW_TOP else False
                self.left_arrow = pictures.get_pygame_image("arrow.png", size, vflip=vflip)
                self.right_arrow = pictures.get_pygame_image("arrow.png", size, hflip=True, vflip=vflip)

                inter = (self.get_rect().width - 2 * self.left_arrow.get_rect().width) // 4

                x0 = int(self.get_rect().left + inter) - x_offset
                x1 = int(self.get_rect().left + 3 * inter + self.left_arrow.get_rect().width) + x_offset

                self.left_arrow_pos = (x0 - self.arrow_offset, y)
                self.right_arrow_pos = (x1 + self.arrow_offset, y)

    def paint(self, screen):
        Background.paint(self, screen)
        screen.blit(self.layout0, self.layout0_pos)
        screen.blit(self.layout1, self.layout1_pos)
        if self.arrow_location != ARROW_HIDDEN:
            screen.blit(self.left_arrow, self.left_arrow_pos)
            screen.blit(self.right_arrow, self.right_arrow_pos)

class ChosenBackground(Background):

    def __init__(self, choices, selected):
        Background.__init__(self, "chosen")
        self.choices = choices
        self.selected = selected
        self.layout = None
        self.layout_pos = None

    def __str__(self):
        return "chosen{}.png".format(self.selected)

    def resize(self, screen):
        Background.resize(self, screen)
        self.text_x = screen.get_rect().centerx
        self.text_y = screen.get_rect().height/7
        if self._need_update:
            size = (self.get_rect().width * 0.6, self.get_rect().height * 0.6)

            self.layout = pictures.get_pygame_image("layout{}.png".format(self.selected), size)

            x = self.layout.get_rect(center=self.get_rect().center).left
            y = int(self.get_rect().top + self.get_rect().height * 0.3)

            self.layout_pos = (x, y)

    def paint(self, screen):
        Background.paint(self, screen)
        screen.blit(self.layout, self.layout_pos)

class CaptureBackground(Background):

    def __init__(self):
        Background.__init__(self, "capture")

class ProcessingBackground(Background):

    def __init__(self):
        Background.__init__(self, "processing")

    def resize(self, screen):
        Background.resize(self, screen)
        self.text_x = screen.get_rect().centerx
        self.text_y = screen.get_rect().height*5/6

class PrintBackground(Background):

    def __init__(self, arrow_location=ARROW_BOTTOM, arrow_offset=0):
        Background.__init__(self, "print")
        self.arrow_location = arrow_location
        self.arrow_offset = arrow_offset
        self.right_arrow = None
        self.right_arrow_pos = None

    def resize(self, screen):
        Background.resize(self, screen)
        self.text_x = screen.get_rect().centerx*3/2
        self.text_y = screen.get_rect().centery
        if self._need_update:
            if self.arrow_location != ARROW_HIDDEN:
                size = (self.get_rect().width * 0.3, self.get_rect().height * 0.3)

                vflip = True if self.arrow_location == ARROW_TOP else False
                self.right_arrow = pictures.get_pygame_image("arrow.png", size, hflip=True, vflip=vflip)

                x = int(self.get_rect().left + self.get_rect().width * 0.75
                        - self.right_arrow.get_rect().width // 2)

                if self.arrow_location == ARROW_TOP:
                    y = self.get_rect().top + 10
                else:
                    y = int(self.get_rect().top + 2 * self.get_rect().height // 3)

                self.right_arrow_pos = (x + self.arrow_offset, y)

    def paint(self, screen):
        Background.paint(self, screen)
        if self.arrow_location != ARROW_HIDDEN:
            screen.blit(self.right_arrow, self.right_arrow_pos)

class FinishedBackground(Background):

    def __init__(self):
        Background.__init__(self, "finished")

    def resize(self, screen):
        Background.resize(self, screen)
        self.text_x = screen.get_rect().centerx
        self.text_y = screen.get_rect().height*3/4

class OopsBackground(Background):

    def __init__(self):
        Background.__init__(self, "oops")

    def resize(self, screen):
        Background.resize(self, screen)
        self.text_x = screen.get_rect().centerx
        self.text_y = screen.get_rect().centery
<commit_msg>Remove useless workaround to have transparent text<commit_after># -*- coding: utf-8 -*-

import os.path as osp

from pibooth import fonts
from pibooth import pictures

from pibooth.language import get_translated_text

ARROW_TOP = 'top'
ARROW_BOTTOM = 'bottom'
ARROW_HIDDEN = 'hidden'

class Background(object):

    def __init__(self, image_name, color=(0, 0, 0), text_color=(255, 255, 255)):
        self._rect = None
        self._need_update = False

        self._background = None
        self._background_color = color
        self._background_image = None

        self._overlay = None
        self._overlay_image = "{0}.png".format(image_name)
        self.text = get_translated_text(image_name)
        self.text_size = 72
        self.text_x = 0
        self.text_y = 0
        self._text_color = text_color

    def __str__(self):
        """Return background final name.
        """
        return self._overlay_image

    def get_rect(self):
        """Return the Rect object of the background image. As aspect ratio
        is kept, the size of the background may be different from the
        screen size.
        """
        return self._overlay.get_rect(center=self._rect.center)

    def set_color(self, color_or_path):
        """Set background color (RGB tuple) or path to an image that used to
        fill the background.

        :param color_or_path: RGB color tuple or image path
        :type color_or_path: tuple or str
        """
        if isinstance(color_or_path, (tuple, list)):
            assert len(color_or_path) == 3, "Length of 3 is required for RGB tuple"
            if color_or_path != self._background_color:
                self._background_color = color_or_path
                self._need_update = True
        else:
            assert osp.isfile(color_or_path), "Invalid image for window background: '{}'".format(color_or_path)
            if color_or_path != self._background_image:
                self._background_image = color_or_path
                self._need_update = True

    def set_text_color(self, color):
        """Set text color (RGB tuple) used to write the texts.
        :param color: RGB color tuple
        :type color: tuple
        """
        assert len(color) == 3, "Length of 3 is required for RGB tuple"
        if color != self._text_color:
            self._text_color = color
            self._need_update = True

    def resize(self, screen):
        """Resize objects to fit to the screen.
        """
        if self._rect != screen.get_rect():
            self._rect = screen.get_rect()
            self.text_size = 72*self._rect.height//400
            self._overlay = pictures.get_pygame_image(self._overlay_image, (self._rect.width, self._rect.height))
            if self._background_image:
                self._background = pictures.get_pygame_image(
                    self._background_image, (self._rect.width, self._rect.height), crop=True)
            self._need_update = True

    def paint(self, screen):
        """Paint and animate the surfaces on the screen.
        """
        if self._background:
            screen.blit(self._background, (0, 0))
        else:
            screen.fill(self._background_color)
        screen.blit(self._overlay, self.get_rect())
        self.write_text(screen)
        self._need_update = False

    def write_text(self, screen):
        """Add the text to the screen
        """
        text_font = pictures.pygame.font.Font(fonts.get_filename("Amatic-Bold"), self.text_size)
        text = text_font.render(self.text, True, self._text_color)
        textrect = text.get_rect()
        textrect.centerx = self.text_x
        textrect.centery = self.text_y
        screen.blit(text, textrect)
        self._need_update = False

class IntroBackground(Background):

    def __init__(self, arrow_location=ARROW_BOTTOM, arrow_offset=0):
        Background.__init__(self, "intro")
        self.arrow_location = arrow_location
        self.arrow_offset = arrow_offset
        self.left_arrow = None
        self.left_arrow_pos = None

    def resize(self, screen):
        Background.resize(self, screen)
        self.text_x = screen.get_rect().centerx/2
        self.text_y = screen.get_rect().centery
        if self._need_update:
            if self.arrow_location != ARROW_HIDDEN:
                size = (self.get_rect().width * 0.3, self.get_rect().height * 0.3)

                vflip = True if self.arrow_location == ARROW_TOP else False
                self.left_arrow = pictures.get_pygame_image("arrow.png", size, vflip=vflip)

                x = int(self.get_rect().left + self.get_rect().width // 4
                        - self.left_arrow.get_rect().width // 2)
                if self.arrow_location == ARROW_TOP:
                    y = self.get_rect().top + 10
                else:
                    y = int(self.get_rect().top + 2 * self.get_rect().height // 3)

                self.left_arrow_pos = (x - self.arrow_offset, y)

    def paint(self, screen):
        Background.paint(self, screen)
        if self.arrow_location != ARROW_HIDDEN:
            screen.blit(self.left_arrow, self.left_arrow_pos)

class IntroWithPrintBackground(IntroBackground):

    def __init__(self, arrow_location=ARROW_BOTTOM, arrow_offset=0):
        IntroBackground.__init__(self, arrow_location, arrow_offset)
        self._overlay_image = "intro_print"
        self.right_arrow = None
        self.right_arrow_pos = None

    def write_text(self, screen):
        IntroBackground.write_text(self, screen)
        text_size = 20*screen.get_rect().height//400
        text_font = pictures.pygame.font.Font(fonts.get_filename("Amatic-Bold"), text_size)
        text_strings = get_translated_text("intro_print").splitlines()
        delta_y = 0
        for text_string in text_strings:
            text = text_font.render(text_string, True, (255, 255, 255))
            textrect = text.get_rect()
            textrect.centerx = self.get_rect().width * 45/100
            textrect.centery = self.get_rect().height * 2/3 + delta_y
            screen.blit(text, textrect)
            delta_y += text.get_height()
        self._need_update = False

    def resize(self, screen):
        IntroBackground.resize(self, screen)
        if self.arrow_location != ARROW_HIDDEN:
            size = (self.get_rect().width * 0.1, self.get_rect().height * 0.1)
            vflip = True if self.arrow_location == ARROW_TOP else False
            self.right_arrow = pictures.get_pygame_image("arrow.png", size, hflip=True, vflip=vflip)

            x = int(self.get_rect().left + self.get_rect().width // 2
                    - self.right_arrow.get_rect().width // 2)
            if self.arrow_location == ARROW_TOP:
                y = self.get_rect().top + 10
            else:
                y = int(self.get_rect().top + 8 * self.get_rect().height // 9)
            self.right_arrow_pos = (x - self.arrow_offset, y)

    def paint(self, screen):
        IntroBackground.paint(self, screen)
        if self.arrow_location != ARROW_HIDDEN:
            screen.blit(self.right_arrow, self.right_arrow_pos)

class ChooseBackground(Background):

    def __init__(self, choices, arrow_location=ARROW_BOTTOM, arrow_offset=0):
        Background.__init__(self, "choose")
        self.arrow_location = arrow_location
        self.arrow_offset = arrow_offset
        self.choices = choices
        self.layout0 = None
        self.layout0_pos = None
        self.layout1 = None
        self.layout1_pos = None
        self.left_arrow = None
        self.left_arrow_pos = None
        self.right_arrow = None
        self.right_arrow_pos = None

    def resize(self, screen):
        Background.resize(self, screen)
        self.text_x = screen.get_rect().centerx
        self.text_y = screen.get_rect().height/7
        if self._need_update:
            size = (self.get_rect().width * 0.6, self.get_rect().height * 0.6)
            self.layout0 = pictures.get_pygame_image("layout{}.png".format(self.choices[0]), size)
            self.layout1 = pictures.get_pygame_image("layout{}.png".format(self.choices[1]), size)

            inter = (self.get_rect().width - 2 * self.layout0.get_rect().width) // 3

            x0 = int(self.get_rect().left + inter)
            x1 = int(self.get_rect().left + 2 * inter + self.layout0.get_rect().width)
            y = int(self.get_rect().top + self.get_rect().height * 0.3)

            self.layout0_pos = (x0, y)
            self.layout1_pos = (x1, y)

            if self.arrow_location != ARROW_HIDDEN:
                if self.arrow_location == ARROW_TOP:
                    y = self.get_rect().top + 5
                    x_offset = 30
                    size = (self.get_rect().width * 0.1, self.layout0_pos[1] - y - 10)
                else:
                    x_offset = 0
                    y = self.layout0_pos[1] + self.layout0.get_rect().height + 5
                    size = (self.get_rect().width * 0.1, self.get_rect().bottom - y - 5)

                vflip = True if self.arrow_location == ARROW_TOP else False
                self.left_arrow = pictures.get_pygame_image("arrow.png", size, vflip=vflip)
                self.right_arrow = pictures.get_pygame_image("arrow.png", size, hflip=True, vflip=vflip)

                inter = (self.get_rect().width - 2 * self.left_arrow.get_rect().width) // 4

                x0 = int(self.get_rect().left + inter) - x_offset
                x1 = int(self.get_rect().left + 3 * inter + self.left_arrow.get_rect().width) + x_offset

                self.left_arrow_pos = (x0 - self.arrow_offset, y)
                self.right_arrow_pos = (x1 + self.arrow_offset, y)

    def paint(self, screen):
        Background.paint(self, screen)
        screen.blit(self.layout0, self.layout0_pos)
        screen.blit(self.layout1, self.layout1_pos)
        if self.arrow_location != ARROW_HIDDEN:
            screen.blit(self.left_arrow, self.left_arrow_pos)
            screen.blit(self.right_arrow, self.right_arrow_pos)

class ChosenBackground(Background):

    def __init__(self, choices, selected):
        Background.__init__(self, "chosen")
        self.choices = choices
        self.selected = selected
        self.layout = None
        self.layout_pos = None

    def __str__(self):
        return "chosen{}.png".format(self.selected)

    def resize(self, screen):
        Background.resize(self, screen)
        self.text_x = screen.get_rect().centerx
        self.text_y = screen.get_rect().height/7
        if self._need_update:
            size = (self.get_rect().width * 0.6, self.get_rect().height * 0.6)

            self.layout = pictures.get_pygame_image("layout{}.png".format(self.selected), size)

            x = self.layout.get_rect(center=self.get_rect().center).left
            y = int(self.get_rect().top + self.get_rect().height * 0.3)

            self.layout_pos = (x, y)

    def paint(self, screen):
        Background.paint(self, screen)
        screen.blit(self.layout, self.layout_pos)

class CaptureBackground(Background):

    def __init__(self):
        Background.__init__(self, "capture")

class ProcessingBackground(Background):

    def __init__(self):
        Background.__init__(self, "processing")

    def resize(self, screen):
        Background.resize(self, screen)
        self.text_x = screen.get_rect().centerx
        self.text_y = screen.get_rect().height*5/6

class PrintBackground(Background):

    def __init__(self, arrow_location=ARROW_BOTTOM, arrow_offset=0):
        Background.__init__(self, "print")
        self.arrow_location = arrow_location
        self.arrow_offset = arrow_offset
        self.right_arrow = None
        self.right_arrow_pos = None

    def resize(self, screen):
        Background.resize(self, screen)
        self.text_x = screen.get_rect().centerx*3/2
        self.text_y = screen.get_rect().centery
        if self._need_update:
            if self.arrow_location != ARROW_HIDDEN:
                size = (self.get_rect().width * 0.3, self.get_rect().height * 0.3)

                vflip = True if self.arrow_location == ARROW_TOP else False
                self.right_arrow = pictures.get_pygame_image("arrow.png", size, hflip=True, vflip=vflip)

                x = int(self.get_rect().left + self.get_rect().width * 0.75
                        - self.right_arrow.get_rect().width // 2)

                if self.arrow_location == ARROW_TOP:
                    y = self.get_rect().top + 10
                else:
                    y = int(self.get_rect().top + 2 * self.get_rect().height // 3)

                self.right_arrow_pos = (x + self.arrow_offset, y)

    def paint(self, screen):
        Background.paint(self, screen)
        if self.arrow_location != ARROW_HIDDEN:
            screen.blit(self.right_arrow, self.right_arrow_pos)

class FinishedBackground(Background):

    def __init__(self):
        Background.__init__(self, "finished")

    def resize(self, screen):
        Background.resize(self, screen)
        self.text_x = screen.get_rect().centerx
        self.text_y = screen.get_rect().height*3/4

class OopsBackground(Background):

    def __init__(self):
        Background.__init__(self, "oops")

    def resize(self, screen):
        Background.resize(self, screen)
        self.text_x = screen.get_rect().centerx
        self.text_y = screen.get_rect().centery
<|endoftext|>


<commit_before># -*- coding: Latin-1 -*-
# Copyright (c) 2014, The MITRE Corporation. All rights reserved.

__version__ = '0.1.0a1'

import argparse
import sys
import os
try:
    import pefile
except ImportError:
    print 'Unable to import pefile'
    try:
        sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
        import pefile
    except Exception as err:
        print 'Unable to import pefile from parent directory.\n'
        print 'Please ensure that pefile resides in the current directory or \
                install pefile.\n'
        print 'ERROR MSG => %s' % err.message
        sys.exit(-1)
from pprint import pprint

from mappings.image_dos_header import IMAGE_DOS_HEADER_MAPPINGS
from mappings.image_optional_header import IMAGE_OPTIONAL_HEADER32_MAPPINGS
from mappings.image_sections import IMAGE_SECTION_HEADER_MAPPINGS
import maec.utils

class PefileParser(object):
    def __init__(self, pe):
        self.pe = pe
        self.pe_file_dictionary = {}
        self.process_pefile_object()

    def perform_mappings(self, struct_instance, object_mapping_dict):
        output_dict = {}
        for key, value in struct_instance.dump_dict().items():
            if key in object_mapping_dict:
                if isinstance(object_mapping_dict[key], str):
                    output_dict[object_mapping_dict[key]] = str(getattr(struct_instance, key))
                elif isinstance(object_mapping_dict[key], dict):
                    mapping_dict = object_mapping_dict[key]
                    # Handle any fields from which the value must be extracted
                    if 'Value' in mapping_dict:
                        child_value = mapping_dict['Value']
                        if child_value:
                            output_dict[mapping_dict['mapping']] = str(child_value.strip())
        return output_dict

    def parse_pe_headers(self):
        headers_dict = {}
        headers_dict['image_dos_header'] = self.perform_mappings(self.pe.DOS_HEADER, IMAGE_DOS_HEADER_MAPPINGS)
        return headers_dict

    def parse_pe_optional_header(self):
        optheaders_dict = {}
        optheaders_dict['image_optional_header'] = self.perform_mappings(self.pe.OPTIONAL_HEADER, IMAGE_OPTIONAL_HEADER32_MAPPINGS)
        return optheaders_dict

    def parse_pe_sections(self):
        sections_dict = {}
        for section in self.pe.sections:
            sections_dict[section.name] = self.perform_mappings(section, IMAGE_SECTION_HEADER_MAPPINGS)
        return sections_dict

    def process_pefile_object(self):
        self.pe_file_dictionary = {'xsi:type':'WindowsExecutableFileObjectType'}
        self.pe_file_dictionary['headers'] = self.parse_pe_headers()
        self.pe_file_dictionary['optional_header'] = self.parse_pe_optional_header()
        self.pe_file_dictionary['sections'] = self.parse_pe_sections()
        pprint(self.pe_file_dictionary)

if __name__ == '__main__':

    # Setup the argument parser
    parser = argparse.ArgumentParser(description="Ero Carrera's Pefile module to MITRE MAEC Translator v" + __version__)
    parser.add_argument("input", help="the name of the input portable executable file to translate to MAEC")
    parser.add_argument("output", help="the name of the MAEC XML to which the output will be written")
    args = parser.parse_args()

    # Open the input file and instantiate pefile.PE
    try:
        pe = pefile.PE(args.input, fast_load=True)
    except pefile.PEFormatError as err:
        print err.message
        sys.exit(-1)
    # Instantiate the pefile parser and parse the pefile object
    parser = PefileParser(pe)
<commit_msg>[pe2m] Evaluating perform_mapping function<commit_after># -*- coding: Latin-1 -*-
# Copyright (c) 2014, The MITRE Corporation. All rights reserved.

__version__ = '0.1.0a1'

import argparse
import sys
import os
try:
    import pefile
except ImportError:
    print 'Unable to import pefile'
    try:
        sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
        import pefile
    except Exception as err:
        print 'Unable to import pefile from parent directory.\n'
        print 'Please ensure that pefile resides in the current directory or \
                install pefile.\n'
        print 'ERROR MSG => %s' % err.message
        sys.exit(-1)
from pprint import pprint

from mappings.image_dos_header import IMAGE_DOS_HEADER_MAPPINGS
from mappings.image_optional_header import IMAGE_OPTIONAL_HEADER32_MAPPINGS
from mappings.image_sections import IMAGE_SECTION_HEADER_MAPPINGS
import maec.utils
from maec.bundle.bundle import Bundle
from maec.bundle.behavior import Behavior
from maec.bundle.malware_action import MalwareAction
from maec.bundle.capability import Capability
from maec.bundle.bundle_reference import BundleReference
from maec.package.analysis import Analysis
from maec.package.malware_subject import MalwareSubject, MalwareSubjectRelationshipList
from maec.package.package import Package
from cybox.utils import Namespace
from cybox.core import AssociatedObjects, AssociatedObject, Object, AssociationType, RelatedObject
from cybox.common.tools import ToolInformation

class PefileToMAEC(object):
    def __init__(self, pefile_parser):
        self.pefile_parser = pefile_parser
        NS = Namespace("http://code.google.com/p/pefile/", "EroCarrera")
        maec.utils.set_id_namespace(NS)
        self.package = Package()
        self.generate_maec()

    def create_object_dict(self, properties_dict):
        object_dict = {'id': maec.utils.idgen.create_id(prefix="object"), 'properties': properties_dict}
        return object_dict

    def populate(self, entry_dict, static_bundle, malware_subject=None):
        object_dict = self.create_object_dict(entry_dict)
        static_bundle.add_object(Object.from_dict(object_dict))

    def generate_analysis(self, static_bundle):
        analysis = Analysis()
        analysis.type = 'triage'
        analysis.method = 'static'
        analysis.add_tool(ToolInformation.from_dict({'id': maec.utils.idgen.create_id(prefix="tool"),
                'vendor': 'Ero Carrera',
                'name': 'pefile'}))
        findings_bundle_reference = []
        if self.bundle_has_content(static_bundle):
            findings_bundle_reference.append(BundleReference.from_dict({'bundle_idref':static_bundle.id}))
        analysis.findings_bundle_reference = findings_bundle_reference
        return analysis

    def bundle_has_content(self, bundle):
        if bundle.actions and len(bundle.actions) > 0:
            return True
        if bundle.objects and len(bundle.objects) > 0:
            return True
        if bundle.behaviors and len(bundle.behaviors) > 0:
            return True
        return False

    def generate_malware_subjects(self):
        entry_dict = self.pefile_parser.pefile_dict
        malware_subject = MalwareSubject()
        entry_dict['id'] = malware_subject
        static_bundle = Bundle(None, False, '4.1', 'static analysis tool output')
        self.populate(entry_dict, static_bundle, malware_subject)
        malware_subject.add_analysis(self.generate_analysis(static_bundle))
        if self.bundle_has_content(static_bundle):
            malware_subject.add_findings_bundle(static_bundle)
        self.package.add_malware_subject(malware_subject)

    def generate_maec(self):
        self.generate_malware_subjects()

class PefileParser(object):
    def __init__(self, pe):
        self.pe = pe
        self.pefile_dict = {}
        self.handle_pefile_object()

    def perform_mapping(self, struct_dict, object_mapping_dict):
        output_dict = {}
        for key, value in struct_dict.items():
            if key in object_mapping_dict.keys():
                if isinstance(object_mapping_dict[key], str):
                    output_dict[object_mapping_dict[key]] = str(struct_dict[key])
                elif isinstance(object_mapping_dict[key], dict):
                    mapping_dict = object_mapping_dict[key]
                    if 'Value' in mapping_dict:
                        child_value = mapping_dict['Value']
                        if child_value:
                            output_dict[mapping_dict['mapping']] = str(child_value.strip())

        pprint(output_dict)
        return output_dict

    def old_perform_mapping(self, struct_instance, object_mapping_dict):
        output_dict = {}
        for key, value in struct_instance.dump_dict().items():
            if key in object_mapping_dict.keys():
                if isinstance(object_mapping_dict[key], str):
                    output_dict[object_mapping_dict[key]] = str(getattr(struct_instance, key))
                elif isinstance(object_mapping_dict[key], dict):
                    mapping_dict = object_mapping_dict[key]
                    # Handle any fields from which the value must be extracted
                    if 'Value' in mapping_dict:
                        child_value = mapping_dict['Value']
                        if child_value:
                            output_dict[mapping_dict['mapping']] = str(child_value.strip())
        return output_dict

    def process_headers(self):
        headers_dict = {}
        struct_dict = self.pe.DOS_HEADER.dump_dict
        headers_dict['dos_header'] = self.perform_mapping(struct_dict(), IMAGE_DOS_HEADER_MAPPINGS)
        return headers_dict

    def handle_pefile_object(self):
        self.pefile_dict = {'xsi:type':'WindowsExecutableFileObjectType'}
        self.pefile_dict['headers'] = {}
        self.pefile_dict['headers'] = self.process_headers()

if __name__ == '__main__':

    # Setup the argument parser
    parser = argparse.ArgumentParser(description="Ero Carrera's Pefile module to MITRE MAEC Translator v" + __version__)
    parser.add_argument("input", help="the name of the input portable executable file to translate to MAEC")
    parser.add_argument("output", help="the name of the MAEC XML to which the output will be written")
    args = parser.parse_args()

    # Open the input file and instantiate pefile.PE
    try:
        pe = pefile.PE(args.input, fast_load=True)
    except pefile.PEFormatError as err:
        print err.message
        sys.exit(-1)
    # Instantiate the pefile parser and parse the pefile object
    parser = PefileParser(pe)
    # Instantiate the MAEC translator and perform the translation
    maec_translator = PefileToMAEC(parser)
    # Output the MAEC Package (generated from pefile output) as XML
    maec_translator.package.to_xml_file(args.output, {"http://code.google.com/p/pefile/":"EroCarrera"})

<|endoftext|>


<commit_before>#!/usr/bin/env python
# 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
#
# Authors:
# - Paul Nilsson, paul.nilsson@cern.ch, 2017

import os
from collections import defaultdict

# from pilot.common.exception import PilotException
from pilot.util.container import execute
from pilot.user.atlas.setup import should_pilot_prepare_asetup, is_user_analysis_job, get_platform, get_asetup, \
    get_asetup_options, is_standard_atlas_job

import logging
logger = logging.getLogger(__name__)

def get_payload_command(job):
    """
    Return the full command for execuring the payload, including the sourcing of all setup files and setting of
    environment variables.

    :param job: job object
    :return: command (string)
    """

    # Should the pilot do the asetup or do the jobPars already contain the information?
    prepareasetup = should_pilot_prepare_asetup(job.get('noExecStrCnv', None), job['jobPars'])

    # Is it a user job or not?
    userjob = is_user_analysis_job(job['transformation'])

    # Get the platform value
    platform = get_platform(job['cmtConfig'])

    # Define the setup for asetup, i.e. including full path to asetup and setting of ATLAS_LOCAL_ROOT_BASE
    asetuppath = get_asetup(asetup=prepareasetup)
    asetupoptions = " "

    if is_standard_atlas_job(job['swRelease']):

        # Normal setup (production and user jobs)
        logger.info("preparing normal production/analysis job setup command")

        cmd = asetuppath
        if prepareasetup:
            options = get_asetup_options(job['swRelease'], job['homepackage'])
            asetupoptions = " " + options + " --platform " + platform

            # Always set the --makeflags option (to prevent asetup from overwriting it)
            asetupoptions += ' --makeflags=\"$MAKEFLAGS\"'

            # Verify that the setup works
            # exitcode, output = timedCommand(cmd, timeout=5 * 60)
            # if exitcode != 0:
            #     if "No release candidates found" in output:
            #         pilotErrorDiag = "No release candidates found"
            #         logger.warning(pilotErrorDiag)
            #         return self.__error.ERR_NORELEASEFOUND, pilotErrorDiag, "", special_setup_cmd, JEM, cmtconfig
            # else:
            #     logger.info("verified setup command")

            cmd += asetupoptions

            if userjob:
                pass
            else:
                # Add Database commands if they are set by the local site
                cmd += os.environ.get('PILOT_DB_LOCAL_SETUP_CMD', '')
                # Add the transform and the job parameters (production jobs)
                if prepareasetup:
                    cmd += ";%s %s" % (job['transformation'], job['jobPars'])
                else:
                    cmd += "; " + job['jobPars']

            cmd = cmd.replace(';;', ';')

    else:  # Generic, non-ATLAS specific jobs, or at least a job with undefined swRelease

        logger.info("generic job (non-ATLAS specific or with undefined swRelease)")

        cmd = ""

    return cmd

def update_job_data(job):
    """
    This function can be used to update/add data to the job object.
    E.g. user specific information can be extracted from other job object fields. In the case of ATLAS, information
    is extracted from the metaData field and added to other job object fields.

    :param job: job object
    :return:
    """

    stageout = "all"

    # handle any error codes
    if 'exeErrorCode' in job['metaData']:
        job['exeErrorCode'] = job['metaData']['exeErrorCode']
        if job['exeErrorCode'] == 0:
            stageout = "all"
        else:
            logger.info('payload failed: exeErrorCode=%d' % job['exeErrorCode'])
            stageout = "log"
    if 'exeErrorDiag' in job['metaData']:
        job['exeErrorDiag'] = job['metaData']['exeErrorDiag']
        if job['exeErrorDiag'] != "":
            logger.warning('payload failed: exeErrorDiag=%s' % job['exeErrorDiag'])

    # determine what should be staged out
    job['stageout'] = stageout  # output and log file or only log file

    # extract the number of events
    job['nEvents'] = get_number_of_events(job['metaData'])

    try:
        work_attributes = parse_jobreport_data(job['metaData'])
    except Exception as e:
        log.warning('failed to parse job report: %s' % e)
    else:
        log.info('work_attributes = %s' % str(work_attributes))

def parse_jobreport_data(job_report):
    """
    Parse a job report and extract relevant fields.

    :param job_report:
    :return:
    """
    work_attributes = {}
    if job_report is None or not any(job_report):
        return work_attributes

    # these are default values for job metrics
    core_count = "undef"
    work_attributes["n_events"] = "undef"
    work_attributes["__db_time"] = "undef"
    work_attributes["__db_data"] = "undef"

    class DictQuery(dict):
        def get(self, path, dst_dict, dst_key):
            keys = path.split("/")
            if len(keys) == 0:
                return
            last_key = keys.pop()
            v = self
            for key in keys:
                if key in v and isinstance(v[key], dict):
                    v = v[key]
                else:
                    return
            if last_key in v:
                dst_dict[dst_key] = v[last_key]
            else:
                return

    if 'ATHENA_PROC_NUMBER' in os.environ:
        work_attributes['core_count'] = os.environ['ATHENA_PROC_NUMBER']
        core_count = os.environ['ATHENA_PROC_NUMBER']

    dq = DictQuery(job_report)
    dq.get("resource/transform/processedEvents", work_attributes, "n_events")
    dq.get("resource/transform/cpuTimeTotal", work_attributes, "cpuConsumptionTime")
    dq.get("resource/machine/node", work_attributes, "node")
    dq.get("resource/machine/model_name", work_attributes, "cpuConsumptionUnit")
    dq.get("resource/dbTimeTotal", work_attributes, "__db_time")
    dq.get("resource/dbDataTotal", work_attributes, "__db_data")
    dq.get("exitCode", work_attributes, "transExitCode")
    dq.get("exitCode", work_attributes, "exeErrorCode")
    dq.get("exitMsg", work_attributes, "exeErrorDiag")
    dq.get("files/input/subfiles", work_attributes, "nInputFiles")

    if 'resource' in job_report and 'executor' in job_report['resource']:
        j = job_report['resource']['executor']
        exc_report = []
        fin_report = defaultdict(int)
        for v in filter(lambda d: 'memory' in d and ('Max' or 'Avg' in d['memory']), j.itervalues()):
            if 'Avg' in v['memory']:
                exc_report.extend(v['memory']['Avg'].items())
            if 'Max' in v['memory']:
                exc_report.extend(v['memory']['Max'].items())
        for x in exc_report:
            fin_report[x[0]] += x[1]
        work_attributes.update(fin_report)

    if 'files' in job_report and 'input' in job_report['files'] and 'subfiles' in job_report['files']['input']:
                work_attributes['nInputFiles'] = len(job_report['files']['input']['subfiles'])

    workdir_size = get_workdir_size()
    work_attributes['jobMetrics'] = 'core_count=%s n_events=%s db_time=%s db_data=%s workdir_size=%s' % \
                                    (core_count,
                                        work_attributes["n_events"],
                                        work_attributes["__db_time"],
                                        work_attributes["__db_data"],
                                        workdir_size)
    del(work_attributes["__db_time"])
    del(work_attributes["__db_data"])

    return work_attributes

def get_workdir_size():
    """
    Tmp function - move later to file_handling

    :return:
    """
    c, o, e = execute('du -s', shell=True)
    if o is not None:
        return o.split()[0]
    return None

def get_executor_dictionary(jobreport_dictionary):
    """
    Extract the 'executor' dictionary from with a job report.

    :param jobreport_dictionary:
    :return: executor_dictionary
    """

    executor_dictionary = {}
    if jobreport_dictionary != {}:

        if 'resource' in jobreport_dictionary:
            resource_dictionary = jobreport_dictionary['resource']
            if 'executor' in resource_dictionary:
                executor_dictionary = resource_dictionary['executor']
            else:
                logger.warning("no such key: executor")
        else:
            logger.warning("no such key: resource")

    return executor_dictionary

def get_number_of_events(jobreport_dictionary):
    """
    Extract the number of events from the job report.

    :param jobreport_dictionary:
    :return:
    """

    nevents = {}  # FORMAT: { format : total_events, .. }
    nmax = 0

    executor_dictionary = get_executor_dictionary(jobreport_dictionary)
    if executor_dictionary != {}:
        for format in executor_dictionary.keys():  # "RAWtoESD", ..
            if 'nevents' in executor_dictionary[format]:
                if format in nevents:
                    print executor_dictionary[format]['nevents']
                    nevents[format] += executor_dictionary[format]['nevents']
                else:
                    nevents[format] = executor_dictionary[format]['nevents']
            else:
                logger.warning("format %s has no such key: nevents" % (format))

    # Now find the largest number of events among the different formats
    if nevents != {}:
        try:
            nmax = max(nevents.values())
        except Exception, e:
            logger.warning("exception caught: %s" % (e))
            nmax = 0
    else:
        logger.warning("did not find the number of events in the job report")
        nmax = 0

    return nmax

def get_db_info(jobreport_dictionary):
    """
    Extract and add up the DB info from the job report.
    This information is reported with the jobMetrics.
    Note: this function adds up the different dbData and dbTime's in the different executor steps. In modern job
    reports this might have been done already by the transform and stored in dbDataTotal and dbTimeTotal.

    :param jobreport_dictionary:
    :return: db_time, db_data
    """

    db_time = 0
    db_data = 0L

    executor_dictionary = get_executor_dictionary(jobreport_dictionary)
    if executor_dictionary != {}:
        for format in executor_dictionary.keys():  # "RAWtoESD", ..
            if 'dbData' in executor_dictionary[format]:
                try:
                    db_data += executor_dictionary[format]['dbData']
                except Exception:
                    pass
            else:
                logger.warning("format %s has no such key: dbData" % (format))
            if 'dbTime' in executor_dictionary[format]:
                try:
                    db_time += executor_dictionary[format]['dbTime']
                except Exception:
                    pass
            else:
                logger.warning("format %s has no such key: dbTime" % (format))

    return db_time, db_data

def get_db_info_str(db_time, db_data):
    """
    Convert db_time, db_data to strings.
    E.g. dbData="105077960", dbTime="251.42".

    :param db_time: time (s)
    :param db_data: long integer
    :return: db_time_s, db_data_s (strings)
    """

    if db_data != 0L:
        db_data_s = "%s" % (db_data)
    else:
        db_data_s = ""
    if db_time != 0:
        db_time_s = "%.2f" % (db_time)
    else:
        db_time_s = ""

    return db_time_s, db_data_s

def get_cpu_times(jobreport_dictionary):
    """
    Extract and add up the total CPU times from the job report.
    E.g. ('s', 5790L, 1.0).

    Note: this function is used with Event Service jobs

    :param jobreport_dictionary:
    :return: cpu_conversion_unit (unit), total_cpu_time, conversion_factor (output consistent with set_time_consumed())
    """

    total_cpu_time = 0L

    executor_dictionary = get_executor_dictionary(jobreport_dictionary)
    if executor_dictionary != {}:
        for format in executor_dictionary.keys():  # "RAWtoESD", ..
            if 'cpuTime' in executor_dictionary[format]:
                try:
                    total_cpu_time += executor_dictionary[format]['cpuTime']
                except Exception:
                    pass
            else:
                logger.warning("format %s has no such key: cpuTime" % (format))

    conversion_factor = 1.0
    cpu_conversion_unit = "s"

    return cpu_conversion_unit, total_cpu_time, conversion_factor

def get_exit_info(jobreport_dictionary):
    """
    Return the exit code (exitCode) and exit message (exitMsg).
    E.g. (0, 'OK').

    :param jobreport_dictionary:
    :return: exit_code, exit_message
    """

    return jobreport_dictionary['exitCode'], jobreport_dictionary['exitMsg']
<commit_msg>Corrected log -> logger<commit_after>#!/usr/bin/env python
# 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
#
# Authors:
# - Paul Nilsson, paul.nilsson@cern.ch, 2017

import os
from collections import defaultdict

# from pilot.common.exception import PilotException
from pilot.util.container import execute
from pilot.user.atlas.setup import should_pilot_prepare_asetup, is_user_analysis_job, get_platform, get_asetup, \
    get_asetup_options, is_standard_atlas_job

import logging
logger = logging.getLogger(__name__)

def get_payload_command(job):
    """
    Return the full command for execuring the payload, including the sourcing of all setup files and setting of
    environment variables.

    :param job: job object
    :return: command (string)
    """

    # Should the pilot do the asetup or do the jobPars already contain the information?
    prepareasetup = should_pilot_prepare_asetup(job.get('noExecStrCnv', None), job['jobPars'])

    # Is it a user job or not?
    userjob = is_user_analysis_job(job['transformation'])

    # Get the platform value
    platform = get_platform(job['cmtConfig'])

    # Define the setup for asetup, i.e. including full path to asetup and setting of ATLAS_LOCAL_ROOT_BASE
    asetuppath = get_asetup(asetup=prepareasetup)
    asetupoptions = " "

    if is_standard_atlas_job(job['swRelease']):

        # Normal setup (production and user jobs)
        logger.info("preparing normal production/analysis job setup command")

        cmd = asetuppath
        if prepareasetup:
            options = get_asetup_options(job['swRelease'], job['homepackage'])
            asetupoptions = " " + options + " --platform " + platform

            # Always set the --makeflags option (to prevent asetup from overwriting it)
            asetupoptions += ' --makeflags=\"$MAKEFLAGS\"'

            # Verify that the setup works
            # exitcode, output = timedCommand(cmd, timeout=5 * 60)
            # if exitcode != 0:
            #     if "No release candidates found" in output:
            #         pilotErrorDiag = "No release candidates found"
            #         logger.warning(pilotErrorDiag)
            #         return self.__error.ERR_NORELEASEFOUND, pilotErrorDiag, "", special_setup_cmd, JEM, cmtconfig
            # else:
            #     logger.info("verified setup command")

            cmd += asetupoptions

            if userjob:
                pass
            else:
                # Add Database commands if they are set by the local site
                cmd += os.environ.get('PILOT_DB_LOCAL_SETUP_CMD', '')
                # Add the transform and the job parameters (production jobs)
                if prepareasetup:
                    cmd += ";%s %s" % (job['transformation'], job['jobPars'])
                else:
                    cmd += "; " + job['jobPars']

            cmd = cmd.replace(';;', ';')

    else:  # Generic, non-ATLAS specific jobs, or at least a job with undefined swRelease

        logger.info("generic job (non-ATLAS specific or with undefined swRelease)")

        cmd = ""

    return cmd

def update_job_data(job):
    """
    This function can be used to update/add data to the job object.
    E.g. user specific information can be extracted from other job object fields. In the case of ATLAS, information
    is extracted from the metaData field and added to other job object fields.

    :param job: job object
    :return:
    """

    stageout = "all"

    # handle any error codes
    if 'exeErrorCode' in job['metaData']:
        job['exeErrorCode'] = job['metaData']['exeErrorCode']
        if job['exeErrorCode'] == 0:
            stageout = "all"
        else:
            logger.info('payload failed: exeErrorCode=%d' % job['exeErrorCode'])
            stageout = "log"
    if 'exeErrorDiag' in job['metaData']:
        job['exeErrorDiag'] = job['metaData']['exeErrorDiag']
        if job['exeErrorDiag'] != "":
            logger.warning('payload failed: exeErrorDiag=%s' % job['exeErrorDiag'])

    # determine what should be staged out
    job['stageout'] = stageout  # output and log file or only log file

    # extract the number of events
    job['nEvents'] = get_number_of_events(job['metaData'])

    try:
        work_attributes = parse_jobreport_data(job['metaData'])
    except Exception as e:
        logger.warning('failed to parse job report: %s' % e)
    else:
        logger.info('work_attributes = %s' % str(work_attributes))

def parse_jobreport_data(job_report):
    """
    Parse a job report and extract relevant fields.

    :param job_report:
    :return:
    """
    work_attributes = {}
    if job_report is None or not any(job_report):
        return work_attributes

    # these are default values for job metrics
    core_count = "undef"
    work_attributes["n_events"] = "undef"
    work_attributes["__db_time"] = "undef"
    work_attributes["__db_data"] = "undef"

    class DictQuery(dict):
        def get(self, path, dst_dict, dst_key):
            keys = path.split("/")
            if len(keys) == 0:
                return
            last_key = keys.pop()
            v = self
            for key in keys:
                if key in v and isinstance(v[key], dict):
                    v = v[key]
                else:
                    return
            if last_key in v:
                dst_dict[dst_key] = v[last_key]
            else:
                return

    if 'ATHENA_PROC_NUMBER' in os.environ:
        work_attributes['core_count'] = os.environ['ATHENA_PROC_NUMBER']
        core_count = os.environ['ATHENA_PROC_NUMBER']

    dq = DictQuery(job_report)
    dq.get("resource/transform/processedEvents", work_attributes, "n_events")
    dq.get("resource/transform/cpuTimeTotal", work_attributes, "cpuConsumptionTime")
    dq.get("resource/machine/node", work_attributes, "node")
    dq.get("resource/machine/model_name", work_attributes, "cpuConsumptionUnit")
    dq.get("resource/dbTimeTotal", work_attributes, "__db_time")
    dq.get("resource/dbDataTotal", work_attributes, "__db_data")
    dq.get("exitCode", work_attributes, "transExitCode")
    dq.get("exitCode", work_attributes, "exeErrorCode")
    dq.get("exitMsg", work_attributes, "exeErrorDiag")
    dq.get("files/input/subfiles", work_attributes, "nInputFiles")

    if 'resource' in job_report and 'executor' in job_report['resource']:
        j = job_report['resource']['executor']
        exc_report = []
        fin_report = defaultdict(int)
        for v in filter(lambda d: 'memory' in d and ('Max' or 'Avg' in d['memory']), j.itervalues()):
            if 'Avg' in v['memory']:
                exc_report.extend(v['memory']['Avg'].items())
            if 'Max' in v['memory']:
                exc_report.extend(v['memory']['Max'].items())
        for x in exc_report:
            fin_report[x[0]] += x[1]
        work_attributes.update(fin_report)

    if 'files' in job_report and 'input' in job_report['files'] and 'subfiles' in job_report['files']['input']:
                work_attributes['nInputFiles'] = len(job_report['files']['input']['subfiles'])

    workdir_size = get_workdir_size()
    work_attributes['jobMetrics'] = 'core_count=%s n_events=%s db_time=%s db_data=%s workdir_size=%s' % \
                                    (core_count,
                                        work_attributes["n_events"],
                                        work_attributes["__db_time"],
                                        work_attributes["__db_data"],
                                        workdir_size)
    del(work_attributes["__db_time"])
    del(work_attributes["__db_data"])

    return work_attributes

def get_workdir_size():
    """
    Tmp function - move later to file_handling

    :return:
    """
    c, o, e = execute('du -s', shell=True)
    if o is not None:
        return o.split()[0]
    return None

def get_executor_dictionary(jobreport_dictionary):
    """
    Extract the 'executor' dictionary from with a job report.

    :param jobreport_dictionary:
    :return: executor_dictionary
    """

    executor_dictionary = {}
    if jobreport_dictionary != {}:

        if 'resource' in jobreport_dictionary:
            resource_dictionary = jobreport_dictionary['resource']
            if 'executor' in resource_dictionary:
                executor_dictionary = resource_dictionary['executor']
            else:
                logger.warning("no such key: executor")
        else:
            logger.warning("no such key: resource")

    return executor_dictionary

def get_number_of_events(jobreport_dictionary):
    """
    Extract the number of events from the job report.

    :param jobreport_dictionary:
    :return:
    """

    nevents = {}  # FORMAT: { format : total_events, .. }
    nmax = 0

    executor_dictionary = get_executor_dictionary(jobreport_dictionary)
    if executor_dictionary != {}:
        for format in executor_dictionary.keys():  # "RAWtoESD", ..
            if 'nevents' in executor_dictionary[format]:
                if format in nevents:
                    print executor_dictionary[format]['nevents']
                    nevents[format] += executor_dictionary[format]['nevents']
                else:
                    nevents[format] = executor_dictionary[format]['nevents']
            else:
                logger.warning("format %s has no such key: nevents" % (format))

    # Now find the largest number of events among the different formats
    if nevents != {}:
        try:
            nmax = max(nevents.values())
        except Exception, e:
            logger.warning("exception caught: %s" % (e))
            nmax = 0
    else:
        logger.warning("did not find the number of events in the job report")
        nmax = 0

    return nmax

def get_db_info(jobreport_dictionary):
    """
    Extract and add up the DB info from the job report.
    This information is reported with the jobMetrics.
    Note: this function adds up the different dbData and dbTime's in the different executor steps. In modern job
    reports this might have been done already by the transform and stored in dbDataTotal and dbTimeTotal.

    :param jobreport_dictionary:
    :return: db_time, db_data
    """

    db_time = 0
    db_data = 0L

    executor_dictionary = get_executor_dictionary(jobreport_dictionary)
    if executor_dictionary != {}:
        for format in executor_dictionary.keys():  # "RAWtoESD", ..
            if 'dbData' in executor_dictionary[format]:
                try:
                    db_data += executor_dictionary[format]['dbData']
                except Exception:
                    pass
            else:
                logger.warning("format %s has no such key: dbData" % (format))
            if 'dbTime' in executor_dictionary[format]:
                try:
                    db_time += executor_dictionary[format]['dbTime']
                except Exception:
                    pass
            else:
                logger.warning("format %s has no such key: dbTime" % (format))

    return db_time, db_data

def get_db_info_str(db_time, db_data):
    """
    Convert db_time, db_data to strings.
    E.g. dbData="105077960", dbTime="251.42".

    :param db_time: time (s)
    :param db_data: long integer
    :return: db_time_s, db_data_s (strings)
    """

    if db_data != 0L:
        db_data_s = "%s" % (db_data)
    else:
        db_data_s = ""
    if db_time != 0:
        db_time_s = "%.2f" % (db_time)
    else:
        db_time_s = ""

    return db_time_s, db_data_s

def get_cpu_times(jobreport_dictionary):
    """
    Extract and add up the total CPU times from the job report.
    E.g. ('s', 5790L, 1.0).

    Note: this function is used with Event Service jobs

    :param jobreport_dictionary:
    :return: cpu_conversion_unit (unit), total_cpu_time, conversion_factor (output consistent with set_time_consumed())
    """

    total_cpu_time = 0L

    executor_dictionary = get_executor_dictionary(jobreport_dictionary)
    if executor_dictionary != {}:
        for format in executor_dictionary.keys():  # "RAWtoESD", ..
            if 'cpuTime' in executor_dictionary[format]:
                try:
                    total_cpu_time += executor_dictionary[format]['cpuTime']
                except Exception:
                    pass
            else:
                logger.warning("format %s has no such key: cpuTime" % (format))

    conversion_factor = 1.0
    cpu_conversion_unit = "s"

    return cpu_conversion_unit, total_cpu_time, conversion_factor

def get_exit_info(jobreport_dictionary):
    """
    Return the exit code (exitCode) and exit message (exitMsg).
    E.g. (0, 'OK').

    :param jobreport_dictionary:
    :return: exit_code, exit_message
    """

    return jobreport_dictionary['exitCode'], jobreport_dictionary['exitMsg']
<|endoftext|>


<commit_before>#    This file is part of WebBox.
#
#    Copyright 2011-2012 Daniel Alexander Smith, Max Van Kleek
#    Copyright 2011-2012 University of Southampton
#
#    WebBox is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    WebBox is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with WebBox.  If not, see <http://www.gnu.org/licenses/>.

import logging, traceback, json
from twisted.web.resource import Resource
from twisted.web.server import NOT_DONE_YET
from webbox.webserver.session import WebBoxSession, ISession
from webbox.exception import ResponseOverride
from mimeparse import quality
from urlparse import parse_qs
try:
    import cjson
    logging.debug("Using CJSON.")
except Exception as e:
    logging.debug("No CJSON, falling back to python json.")

class BaseHandler(Resource):
    """ Add/remove boxes, add/remove users, change config. """

    base_path=None  # Override me , e.g., 'auth'

    subhandlers = {
        # e.g.,
        # 'login': {
        #     'methods': ['POST'],
        #     'require_auth': False,
        #     'require_token': False,
        #     'handler': self.login,
        #     'content-type':'text/plain' # optional
        # }
    }

    def __init__(self, webserver, base_path=None, register=True):
        Resource.__init__(self)
        self.webserver = webserver

        if base_path is not None:
            self.base_path = base_path

        self.isLeaf = True # stops twisted from seeking children resources from me
        if register:
            logging.debug("Putting child " + (self.base_path or ''))
            webserver.root.putChild(self.base_path, self) # register path with webserver

    def _matches_request(self, request, subhandler):
        path_fields = request.path.split("/")
        sub_path = '/'.join(path_fields[2:]) if len(path_fields) >= 3 else ''

        #logging.debug('sub_path {0} {1}'.format(sub_path, subhandler['prefix']))

        # if the subhandler supports content negotiation, then determine the best one
        assert subhandler['accept'], 'No accept clause in subhandler %s ' % subhandler['prefix']

        if not (subhandler["prefix"] == sub_path or subhandler["prefix"] == '*'):
            # logging.debug("__PREFIX mismatch " + sub_path + "  "  + subhandler["prefix"])
            #logging.debug('prefix mismatch')
            return False
        if self._get_best_content_type_match_score(request,subhandler) <= 0:
            # logging.debug("__NOT content type match " + self._get_best_content_type_match_score(request,subhandler))
            #logging.debug('content type mismatch ')
            return False
        if not request.method in subhandler["methods"]:
            # logging.debug("__NOT in subhandler " + request.method)
            #logging.debug('method type mismatch ' + request.method + ' ' + repr(subhandler["methods"]))
            return False
        logging.debug("Handler match where path={0} and method={1}".format(sub_path, request.method))
        return True

    def _get_best_content_type_match_score(self,request,subhandler):
        request_accept = request.getHeader('Accept') or '*/*'
        return max(map(lambda handler_mimetype:quality(handler_mimetype,request_accept), subhandler["accept"]))

    def _matches_auth_requirements(self, request, subhandler):
        session = self.get_session(request)
        if subhandler['require_auth'] and not session.is_authenticated:
            return False
        # @TODO
        # if subhandler['require_token'] and not true:
        #    raise ForbiddenResource()
        return True

    def _get_arg(self,request,argname):
        if request.method == 'GET':
            return request.args.get(argname) and request.args[argname][0]
        if request.method in ['POST', 'PUT', 'DELETE', 'COPY', 'MOVE']:
            post_args = self.get_post_args(request)
            return post_args.get(argname) and post_args[argname][0]
        return None

    def get_token(self,request):
        tid = self._get_arg(request,'token')
        return self.webserver.tokens.get(tid) if tid else None

    def get_origin(self,request):
        return request.getHeader('origin')

    ## new in apps-refactor version - requires box name to be
    ## provided in the request.
    ## what we _used_ to do is require
    ##   "apps/enriches/get_next_round" !== "box/url" -> fail.
    ## ??
    def get_request_box(self,request):
        return self._get_arg(request,'box')
    def get_request_app(self,request):
        return self._get_arg(request,'app')

    # revision to protocol
    def _matches_token_requirements(self, request, subhandler):
        if not subhandler['require_token']: return True
        token,boxid,appid = self.get_token(request), self.get_request_box(request), self.get_request_app(request)
        return token and token.verify(boxid, appid, self.get_origin(request))

    def get_session(self,request):
        session = request.getSession()
        # persists for life of a session (based on the cookie set by the above)
        wbSession = session.getComponent(ISession)
        if not wbSession:
            wbSession = WebBoxSession(session, self.webserver)
            session.setComponent(ISession, wbSession)
        return wbSession

    def render(self, request):
        """ Twisted resource handler."""
        logging.debug("Calling base render() - " + '/'.join(request.path.split("/")) + " " + repr( self.__class__)  + ' _ subhandlers::' + repr(len(self.subhandlers)))
        try:
            self.set_cors_headers(request)
            matching_handlers = filter(lambda h: self._matches_request(request,h), self.subhandlers)
            logging.debug('Matching handlers %d' % len(matching_handlers))
            matching_handlers.sort(key=lambda h: self._get_best_content_type_match_score(request,h),reverse=True)
            matching_auth_hs = filter(lambda h: self._matches_auth_requirements(request,h), matching_handlers)
            logging.debug('Post-auth matching handlers %d' % len(matching_auth_hs))
            matching_token_hs = filter(lambda h: self._matches_token_requirements(request,h), matching_auth_hs)
            logging.debug('Post-token matching handlers %d' % len(matching_token_hs))

            if matching_token_hs:
                subhandler = matching_token_hs[0]
                logging.debug('Using handler %s' % self.__class__.__name__ + " " + matching_token_hs[0]["prefix"])
                if subhandler['content-type']:
                    request.setHeader('Content-Type', subhandler['content-type'])
                subhandler['handler'](self,request)
                return NOT_DONE_YET
            logging.debug('Returning not found ')
            self.return_not_found(request)
            return NOT_DONE_YET
        except ResponseOverride as roe:
            self._respond(request,roe.status,roe.reason)
            return NOT_DONE_YET
        except Exception as e:
            logging.debug("Error in AdminHandler.render(), returning 500: %s, exception is: %s" % (str(e), traceback.format_exc()))
            self.return_internal_error(request)
            return NOT_DONE_YET
        # never get here
        pass

    def _respond(self, request, code, message, additional_data=None):
        response = {"message": message, "code": code}
        if additional_data:
            response.update(additional_data)
        try:
            responsejson = cjson.encode(response)
            logging.debug("Encoding response with cjson")
        except Exception as e:
            responsejson = json.dumps(response)
            logging.debug("Encoding response with python json")
        request.setResponseCode(code, message=message)
        request.setHeader("Content-Type", "application/json")
        request.setHeader("Content-Length", len(responsejson))
        request.write(responsejson)
        request.finish()
        logging.debug(' just called request.finish() with code %d ' % code)

    def return_ok(self,request,data=None):
        self._respond(request, 200, "OK", data)
    def return_created(self,request,data=None):
        self._respond(request, 201, "Created",data)
    def return_not_found(self,request):
        self._respond(request, 404, "Not Found")
    def return_forbidden(self,request):
        self._respond(request, 403, "Forbidden")
    def return_unauthorized(self,request):
        self._respond(request, 401, "Unauthorized")
    def return_unsupported_method(self,request):
        self._respond(request, 405, "Method Not Allowed")
    def return_internal_error(self,request):
        self._respond(request, 500, "Internal Server Error")
    def return_bad_request(self,request,description=None):
        self._respond(request, 400, 'Bad Request', {"description": description} if description else None)
    def return_obsolete(self,request,data=None):
        self._respond(request, 409, 'Obsolete', data)

    ## allowed for cors
    def get_cors_methods(self, request):
        # default set of allowed methods
        return ("POST", "GET", "PUT", "HEAD", "OPTIONS")
    def get_cors_origin(self, request):
        # default set of allowed origin hosts
        origin = request.getHeader("origin")
        if origin is None:
            return ()
        return (request.getHeader("origin"),) # ("*",)
    def get_cors_headers(self, request):
        # default set
        return ("Content-Type", "origin", "accept", "Depth", "User-Agent", "X-File-Size", "X-Requested-With", "Cookie", "Set-Cookie", "If-Modified-Since","X-File-Name", "Cache-Control")

    def get_post_args(self,request):
        request.content.seek(0) # dan's fault.
        return parse_qs(request.content.read())

    def set_cors_headers(self,request):
        if self.get_cors_origin(request):
            # logging.debug('setting allow origin {0}'.format(' '.join(self.get_cors_origin(request))))
            request.setHeader("Access-Control-Allow-Origin", ' '.join(self.get_cors_origin(request)) )
            request.setHeader("Access-Control-Allow-Methods", ', '.join( self.get_cors_methods(request)))
            request.setHeader("Access-Control-Allow-Headers", ', '.join( self.get_cors_headers(request)) )
            request.setHeader("Access-Control-Allow-Credentials", 'true')
        else:
            #logging.debug('local request - skipping cors')
            pass

<commit_msg>Don't write to the connection if it has already disconnected<commit_after>#    This file is part of WebBox.
#
#    Copyright 2011-2012 Daniel Alexander Smith, Max Van Kleek
#    Copyright 2011-2012 University of Southampton
#
#    WebBox is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    WebBox is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with WebBox.  If not, see <http://www.gnu.org/licenses/>.

import logging, traceback, json
from twisted.web.resource import Resource
from twisted.web.server import NOT_DONE_YET
from webbox.webserver.session import WebBoxSession, ISession
from webbox.exception import ResponseOverride
from mimeparse import quality
from urlparse import parse_qs
try:
    import cjson
    logging.debug("Using CJSON.")
except Exception as e:
    logging.debug("No CJSON, falling back to python json.")

class BaseHandler(Resource):
    """ Add/remove boxes, add/remove users, change config. """

    base_path=None  # Override me , e.g., 'auth'

    subhandlers = {
        # e.g.,
        # 'login': {
        #     'methods': ['POST'],
        #     'require_auth': False,
        #     'require_token': False,
        #     'handler': self.login,
        #     'content-type':'text/plain' # optional
        # }
    }

    def __init__(self, webserver, base_path=None, register=True):
        Resource.__init__(self)
        self.webserver = webserver

        if base_path is not None:
            self.base_path = base_path

        self.isLeaf = True # stops twisted from seeking children resources from me
        if register:
            logging.debug("Putting child " + (self.base_path or ''))
            webserver.root.putChild(self.base_path, self) # register path with webserver

    def _matches_request(self, request, subhandler):
        path_fields = request.path.split("/")
        sub_path = '/'.join(path_fields[2:]) if len(path_fields) >= 3 else ''

        #logging.debug('sub_path {0} {1}'.format(sub_path, subhandler['prefix']))

        # if the subhandler supports content negotiation, then determine the best one
        assert subhandler['accept'], 'No accept clause in subhandler %s ' % subhandler['prefix']

        if not (subhandler["prefix"] == sub_path or subhandler["prefix"] == '*'):
            # logging.debug("__PREFIX mismatch " + sub_path + "  "  + subhandler["prefix"])
            #logging.debug('prefix mismatch')
            return False
        if self._get_best_content_type_match_score(request,subhandler) <= 0:
            # logging.debug("__NOT content type match " + self._get_best_content_type_match_score(request,subhandler))
            #logging.debug('content type mismatch ')
            return False
        if not request.method in subhandler["methods"]:
            # logging.debug("__NOT in subhandler " + request.method)
            #logging.debug('method type mismatch ' + request.method + ' ' + repr(subhandler["methods"]))
            return False
        logging.debug("Handler match where path={0} and method={1}".format(sub_path, request.method))
        return True

    def _get_best_content_type_match_score(self,request,subhandler):
        request_accept = request.getHeader('Accept') or '*/*'
        return max(map(lambda handler_mimetype:quality(handler_mimetype,request_accept), subhandler["accept"]))

    def _matches_auth_requirements(self, request, subhandler):
        session = self.get_session(request)
        if subhandler['require_auth'] and not session.is_authenticated:
            return False
        # @TODO
        # if subhandler['require_token'] and not true:
        #    raise ForbiddenResource()
        return True

    def _get_arg(self,request,argname):
        if request.method == 'GET':
            return request.args.get(argname) and request.args[argname][0]
        if request.method in ['POST', 'PUT', 'DELETE', 'COPY', 'MOVE']:
            post_args = self.get_post_args(request)
            return post_args.get(argname) and post_args[argname][0]
        return None

    def get_token(self,request):
        tid = self._get_arg(request,'token')
        return self.webserver.tokens.get(tid) if tid else None

    def get_origin(self,request):
        return request.getHeader('origin')

    ## new in apps-refactor version - requires box name to be
    ## provided in the request.
    ## what we _used_ to do is require
    ##   "apps/enriches/get_next_round" !== "box/url" -> fail.
    ## ??
    def get_request_box(self,request):
        return self._get_arg(request,'box')
    def get_request_app(self,request):
        return self._get_arg(request,'app')

    # revision to protocol
    def _matches_token_requirements(self, request, subhandler):
        if not subhandler['require_token']: return True
        token,boxid,appid = self.get_token(request), self.get_request_box(request), self.get_request_app(request)
        return token and token.verify(boxid, appid, self.get_origin(request))

    def get_session(self,request):
        session = request.getSession()
        # persists for life of a session (based on the cookie set by the above)
        wbSession = session.getComponent(ISession)
        if not wbSession:
            wbSession = WebBoxSession(session, self.webserver)
            session.setComponent(ISession, wbSession)
        return wbSession

    def render(self, request):
        """ Twisted resource handler."""
        logging.debug("Calling base render() - " + '/'.join(request.path.split("/")) + " " + repr( self.__class__)  + ' _ subhandlers::' + repr(len(self.subhandlers)))
        try:
            self.set_cors_headers(request)
            matching_handlers = filter(lambda h: self._matches_request(request,h), self.subhandlers)
            logging.debug('Matching handlers %d' % len(matching_handlers))
            matching_handlers.sort(key=lambda h: self._get_best_content_type_match_score(request,h),reverse=True)
            matching_auth_hs = filter(lambda h: self._matches_auth_requirements(request,h), matching_handlers)
            logging.debug('Post-auth matching handlers %d' % len(matching_auth_hs))
            matching_token_hs = filter(lambda h: self._matches_token_requirements(request,h), matching_auth_hs)
            logging.debug('Post-token matching handlers %d' % len(matching_token_hs))

            if matching_token_hs:
                subhandler = matching_token_hs[0]
                logging.debug('Using handler %s' % self.__class__.__name__ + " " + matching_token_hs[0]["prefix"])
                if subhandler['content-type']:
                    request.setHeader('Content-Type', subhandler['content-type'])
                subhandler['handler'](self,request)
                return NOT_DONE_YET
            logging.debug('Returning not found ')
            self.return_not_found(request)
            return NOT_DONE_YET
        except ResponseOverride as roe:
            self._respond(request,roe.status,roe.reason)
            return NOT_DONE_YET
        except Exception as e:
            logging.debug("Error in AdminHandler.render(), returning 500: %s, exception is: %s" % (str(e), traceback.format_exc()))
            self.return_internal_error(request)
            return NOT_DONE_YET
        # never get here
        pass

    def _respond(self, request, code, message, additional_data=None):
        response = {"message": message, "code": code}
        if additional_data:
            response.update(additional_data)
        try:
            responsejson = cjson.encode(response)
            logging.debug("Encoding response with cjson")
        except Exception as e:
            responsejson = json.dumps(response)
            logging.debug("Encoding response with python json")

        if not request._disconnected:
            request.setResponseCode(code, message=message)
            request.setHeader("Content-Type", "application/json")
            request.setHeader("Content-Length", len(responsejson))
            request.write(responsejson)
            request.finish()
            logging.debug(' just called request.finish() with code %d ' % code)
        else:
            logging.debug(' didnt call request.finish(), because it was already disconnected')

    def return_ok(self,request,data=None):
        self._respond(request, 200, "OK", data)
    def return_created(self,request,data=None):
        self._respond(request, 201, "Created",data)
    def return_not_found(self,request):
        self._respond(request, 404, "Not Found")
    def return_forbidden(self,request):
        self._respond(request, 403, "Forbidden")
    def return_unauthorized(self,request):
        self._respond(request, 401, "Unauthorized")
    def return_unsupported_method(self,request):
        self._respond(request, 405, "Method Not Allowed")
    def return_internal_error(self,request):
        self._respond(request, 500, "Internal Server Error")
    def return_bad_request(self,request,description=None):
        self._respond(request, 400, 'Bad Request', {"description": description} if description else None)
    def return_obsolete(self,request,data=None):
        self._respond(request, 409, 'Obsolete', data)

    ## allowed for cors
    def get_cors_methods(self, request):
        # default set of allowed methods
        return ("POST", "GET", "PUT", "HEAD", "OPTIONS")
    def get_cors_origin(self, request):
        # default set of allowed origin hosts
        origin = request.getHeader("origin")
        if origin is None:
            return ()
        return (request.getHeader("origin"),) # ("*",)
    def get_cors_headers(self, request):
        # default set
        return ("Content-Type", "origin", "accept", "Depth", "User-Agent", "X-File-Size", "X-Requested-With", "Cookie", "Set-Cookie", "If-Modified-Since","X-File-Name", "Cache-Control")

    def get_post_args(self,request):
        request.content.seek(0) # dan's fault.
        return parse_qs(request.content.read())

    def set_cors_headers(self,request):
        if self.get_cors_origin(request):
            # logging.debug('setting allow origin {0}'.format(' '.join(self.get_cors_origin(request))))
            request.setHeader("Access-Control-Allow-Origin", ' '.join(self.get_cors_origin(request)) )
            request.setHeader("Access-Control-Allow-Methods", ', '.join( self.get_cors_methods(request)))
            request.setHeader("Access-Control-Allow-Headers", ', '.join( self.get_cors_headers(request)) )
            request.setHeader("Access-Control-Allow-Credentials", 'true')
        else:
            #logging.debug('local request - skipping cors')
            pass

<|endoftext|>