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