1c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)# Copyright 2013 The Chromium Authors. All rights reserved.
2c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)# Use of this source code is governed by a BSD-style license that can be
3c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)# found in the LICENSE file.
4c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
5c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)import hashlib
6c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)import os
7c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
8c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
9c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)def CallAndRecordIfStale(
1023730a6e56a168d1879203e4b3819bb36e3d8f1fTorne (Richard Coles)    function, record_path=None, input_paths=None, input_strings=None,
1123730a6e56a168d1879203e4b3819bb36e3d8f1fTorne (Richard Coles)    force=False):
12c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)  """Calls function if the md5sum of the input paths/strings has changed.
13c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
14c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)  The md5sum of the inputs is compared with the one stored in record_path. If
15c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)  this has changed (or the record doesn't exist), function will be called and
16c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)  the new md5sum will be recorded.
17c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
18c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)  If force is True, the function will be called regardless of whether the
19c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)  md5sum is out of date.
20c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)  """
2123730a6e56a168d1879203e4b3819bb36e3d8f1fTorne (Richard Coles)  if not input_paths:
2223730a6e56a168d1879203e4b3819bb36e3d8f1fTorne (Richard Coles)    input_paths = []
2323730a6e56a168d1879203e4b3819bb36e3d8f1fTorne (Richard Coles)  if not input_strings:
2423730a6e56a168d1879203e4b3819bb36e3d8f1fTorne (Richard Coles)    input_strings = []
25c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)  md5_checker = _Md5Checker(
26c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)      record_path=record_path,
27c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)      input_paths=input_paths,
28c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)      input_strings=input_strings)
29c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)  if force or md5_checker.IsStale():
30c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)    function()
31c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)    md5_checker.Write()
32c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
33c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
34c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)def _UpdateMd5ForFile(md5, path, block_size=2**16):
35c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)  with open(path, 'rb') as infile:
36c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)    while True:
37c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)      data = infile.read(block_size)
38c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)      if not data:
39c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)        break
40c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)      md5.update(data)
41c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
42c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
43c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)def _UpdateMd5ForDirectory(md5, dir_path):
44c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)  for root, _, files in os.walk(dir_path):
45c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)    for f in files:
46c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)      _UpdateMd5ForFile(md5, os.path.join(root, f))
47c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
48c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
49c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)def _UpdateMd5ForPath(md5, path):
50c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)  if os.path.isdir(path):
51c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)    _UpdateMd5ForDirectory(md5, path)
52c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)  else:
53c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)    _UpdateMd5ForFile(md5, path)
54c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
55c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
56c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)class _Md5Checker(object):
5723730a6e56a168d1879203e4b3819bb36e3d8f1fTorne (Richard Coles)  def __init__(self, record_path=None, input_paths=None, input_strings=None):
5823730a6e56a168d1879203e4b3819bb36e3d8f1fTorne (Richard Coles)    if not input_paths:
5923730a6e56a168d1879203e4b3819bb36e3d8f1fTorne (Richard Coles)      input_paths = []
6023730a6e56a168d1879203e4b3819bb36e3d8f1fTorne (Richard Coles)    if not input_strings:
6123730a6e56a168d1879203e4b3819bb36e3d8f1fTorne (Richard Coles)      input_strings = []
6223730a6e56a168d1879203e4b3819bb36e3d8f1fTorne (Richard Coles)
63c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)    assert record_path.endswith('.stamp'), (
64c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)        'record paths must end in \'.stamp\' so that they are easy to find '
65c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)        'and delete')
66c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
67c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)    self.record_path = record_path
68c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
69c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)    md5 = hashlib.md5()
70c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)    for i in sorted(input_paths):
71c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)      _UpdateMd5ForPath(md5, i)
72c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)    for s in input_strings:
73c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)      md5.update(s)
74c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)    self.new_digest = md5.hexdigest()
75c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
76c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)    self.old_digest = ''
77c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)    if os.path.exists(self.record_path):
78c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)      with open(self.record_path, 'r') as old_record:
79c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)        self.old_digest = old_record.read()
80c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
81c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)  def IsStale(self):
82c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)    return self.old_digest != self.new_digest
83c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)
84c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)  def Write(self):
85c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)    with open(self.record_path, 'w') as new_record:
86c2e0dbddbe15c98d52c4786dac06cb8952a8ae6dTorne (Richard Coles)      new_record.write(self.new_digest)
87