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

Source Code for Module lib.api.screenshot

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