[AFL] Add update script.

third_party/afl/update.py

Index: third_party/afl/update.py
diff --git a/third_party/afl/update.py b/third_party/afl/update.py
new file mode 100755
index 0000000000000000000000000000000000000000..0dd8df90a34d277b0ee4d9b0edf92eff29f5466d
--- /dev/null
+++ b/third_party/afl/update.py
@@ -0,0 +1,181 @@
+#! /usr/bin/env python
+# Copyright (c) 2016 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Script for updating AFL. Also updates AFL version in README.chromium.
+"""
+
+import argparse
+import cStringIO
+import datetime
+import os
+import re
+import sys
+import tarfile
+import urllib2
+
+
+VERSION_REGEX = r'(?P<version>([0-9]*[.])?[0-9]+b)'
+PATH_REGEX = r'(afl-)' + VERSION_REGEX
+
+
+class ChromiumReadme(object):
+  """Class that handles reading from and updating the README.chromium"""
+
+  README_FILE_PATH = 'third_party/afl/README.chromium'
+  README_VERSION_REGEX = r'Version: ' + VERSION_REGEX
+
+  def __init__(self):
+    """
+    Inits the ChromiumReadme.
+    """
+    with open(self.README_FILE_PATH) as readme_file_handle:
+      self.readme_contents = readme_file_handle.read()
+
+  def get_current_version(self):
+    """
+    Get the current version of AFL according to the README.chromium
+    """
+    match = re.search(self.README_VERSION_REGEX, self.readme_contents)
+    if not match:
+      raise Exception('Could not determine current AFL version')
+    return match.groupdict()['version']
+
+  def update(self, new_version):
+    """
+    Update the readme to reflect the new version that has been downloaded.
+    """
+    new_readme = self.readme_contents
+    subsitutions = [(VERSION_REGEX, new_version),  # Update the version.
+                    (r'Date: .*',
+                     'Date: ' + datetime.date.today().strftime("%B %d, %Y")),
+                    # Update the Local Modifications.
+                    (PATH_REGEX + r'/', 'afl-' + new_version + '/')]
+
+    for regex, replacement in subsitutions:
+      new_readme = re.subn(regex, replacement, new_readme, 1)[0]
+
+    self.readme_contents = new_readme
+    with open(self.README_FILE_PATH, 'w+') as readme_file_handle:
+      readme_file_handle.write(self.readme_contents)
+
+
+class AflTarball(object):
+  """
+  Class that handles the afl-latest.tgz tarball.
+  """
+  # Regexes that match files that we don't want to extract.
+  # Note that you should add these removals to "Local Modifications" in
+  # the README.chromium.
+  UNWANTED_FILE_REGEX = '|'.join([
+      r'(.*\.elf)',  # presubmit complains these aren't marked executable.
+      r'(.*others/elf)',  # We don't need this if we have no elfs.
+      # checkdeps complains about #includes.
+      r'(.*afl-llvm-pass\.so\.cc)',
+      r'(.*argv.*)',  # Delete the demo's directory as well.
+
+      r'(.*dictionaries.*)',  # Including these make builds fail.
+  ])
+  AFL_SRC_DIR = 'third_party/afl/src'
+
+  def __init__(self, version):
+    """
+    Init this AFL tarball.
+    """
+    release_name = 'afl-{0}'.format(version)
+    filename = '{0}.tgz'.format(release_name)
+
+    # Note: lcamtuf.coredump.cx does not support TLS connections. The "http://"
+    # protocol is intentional.
+    self.url = "http://lcamtuf.coredump.cx/afl/releases/{0}".format(filename)
+    self.tarball = None
+    self.real_version = version if version != 'latest' else None
+
+  def download(self):
+    """Download the tarball version from
+    http://lcamtuf.coredump.cx/afl/releases/
+    """
+    tarball_contents = urllib2.urlopen(self.url).read()
+    tarball_file = cStringIO.StringIO(tarball_contents)
+    self.tarball = tarfile.open(fileobj=tarball_file, mode="r:gz")
+    if self.real_version is None:
+      regex_match = re.search(VERSION_REGEX, self.tarball.members[0].path)
+      self.real_version = regex_match.groupdict()['version']
+
+  def extract(self):
+    """
+    Extract the files and folders from the tarball we have downloaded while
+    skipping unwanted ones.
+    """
+
+    for member in self.tarball.getmembers():
+      member.path = re.sub(PATH_REGEX, self.AFL_SRC_DIR, member.path)
+      if re.match(self.UNWANTED_FILE_REGEX, member.path):
+        print 'skipping unwanted file: {0}'.format(member.path)
+        continue
+      self.tarball.extract(member)
+
+
+def version_to_float(version):
+  """
+  Convert version string to float.
+  """
+  if version.endswith('b'):
+    return float(version[:-1])
+
+  return float(version)
+
+
+def update_afl(new_version):
+  """
+  Update this version of AFL to newer version, new_version.
+  """
+  readme = ChromiumReadme()
+  old_version = readme.get_current_version()
+  if new_version != 'latest':
+    new_float = version_to_float(new_version)
+    assert version_to_float(old_version) < new_float, (
+        'Trying to update from version {0} to {1}'.format(old_version,
+                                                          new_version))
+
+  # Extract the tarball.
+  tarball = AflTarball(new_version)
+  tarball.download()
+  current_less_than_new = version_to_float(old_version) < version_to_float(
+      tarball.real_version)
+  assert current_less_than_new, (
+      'Trying to update from version {0} to {1}'.format(old_version,
+                                                        tarball.real_version))
+
+  tarball.extract()
+
+  readme.update(tarball.real_version)
+
+
+def main():
+  """
+  Update AFL if possible.
+  """
+  parser = argparse.ArgumentParser('Update AFL.')
+  parser.add_argument('version', metavar='version', default='latest', nargs='?',
+                      help='(optional) Version to update AFL to.')
+
+  args = parser.parse_args()
+  version = args.version
+  if version != 'latest' and not version.endswith('b'):
+    version += 'b'
+
+  in_correct_directory = (os.path.basename(os.getcwd()) == 'src' and
+                          os.path.exists('third_party'))
+
+  assert in_correct_directory, (
+      '{0} must be run from the repo\'s root'.format(sys.argv[0]))
+
+  update_afl(version)
+  print ("Run git diff third_party/afl/src/docs/ChangeLog to see changes to AFL"
+         " since the last roll")
+
+
+if __name__ == '__main__':
+  main()