lib.api.screenshot

Package lib :: Package api :: Module screenshot |
[hide private]
[frames] | no frames]

# Source Code for Module lib.api.screenshot

     1  # Copyright (C) 2010-2013 Claudio Guarnieri.
     2  # Copyright (C) 2014-2016 Cuckoo Foundation.
     3  # This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
     4  # See the file 'docs/LICENSE' for copying permission.
     5
     6  import math
     7
     8  try:
     9      import ImageChops
    10      import ImageGrab
    11      import ImageDraw
    12      HAVE_PIL = True
    13  except:
    14      try:
    15          from PIL import ImageChops
    16          from PIL import ImageGrab
    17          from PIL import ImageDraw
    18          HAVE_PIL = True
    19      except:
    20          HAVE_PIL = False
    21
    22 -class Screenshot:
    23      """Get screenshots."""
    24
    25 -    def _draw_rectangle(self, img, xy):
    26          """Draw a black rectangle.
    27          @param img: PIL Image object
    28          @param xy: Coordinates as refined in PIL rectangle() doc
    29          @return: Image with black rectangle
    30          """
    31          dr = ImageDraw.Draw(img)
    32          dr.rectangle(xy, fill="black", outline="black")
    33          return img
    34
    35 -    def have_pil(self):
    36          """Is Python Image Library installed?
    37          @return: installed status.
    38          """
    39          return HAVE_PIL
    40
    41 -    def equal(self, img1, img2, skip_area=None):
    42          """Compares two screenshots using Root-Mean-Square Difference (RMS).
    43          @param img1: screenshot to compare.
    44          @param img2: screenshot to compare.
    45          @return: equal status.
    46          """
    47          if not HAVE_PIL:
    48              return None
    49
    50          # Trick to avoid getting a lot of screen shots only because the time in the windows
    51          # clock is changed.
    52          # We draw a black rectangle on the coordinates where the clock is locates, and then
    53          # run the comparison.
    54          # NOTE: the coordinates are changing with VM screen resolution.
    55          if skip_area:
    56              # Copying objects to draw in another object.
    57              img1 = img1.copy()
    58              img2 = img2.copy()
    59              # Draw a rectangle to cover windows clock.
    60              for img in (img1, img2):
    61                  self._draw_rectangle(img, skip_area)
    62
    63          # To get a measure of how similar two images are, we use
    64          # root-mean-square (RMS). If the images are exactly identical,
    65          # this value is zero.
    66          diff = ImageChops.difference(img1, img2)
    67          h = diff.histogram()
    68          sq = (value * ((idx % 256)**2) for idx, value in enumerate(h))
    69          sum_of_squares = sum(sq)
    70          rms = math.sqrt(sum_of_squares/float(img1.size[0] * img1.size[1]))
    71
    72          # Might need to tweak the threshold.
    73          return rms < 8
    74
    75 -    def take(self):
    76          """Take a screenshot.
    77          @return: screenshot or None.
    78          """
    79          if not HAVE_PIL:
    80              return None
    81
    82          return ImageGrab.grab()
    83