net/data/verify_signed_data_unittest/annotate_test_data.py

   1 #!/usr/bin/python
   2 # Copyright (c) 2015 The Chromium Authors. All rights reserved.
   3 # Use of this source code is governed by a BSD-style license that can be
   4 # found in the LICENSE file.
   5
   6 """This script is called without any arguments to re-format all of the *.pem
   7 files in the script's parent directory.
   8
   9 The main formatting change is to run "openssl asn1parse" for each of the PEM
  10 block sections (except for DATA), and add that output to the comment.
  11
  12 Refer to the README file for more information.
  13 """
  14
  15 import glob
  16 import os
  17 import re
  18 import base64
  19 import subprocess
  20
  21
  22 def Transform(file_data):
  23   """Returns a transformed (formatted) version of file_data"""
  24
  25   result = ''
  26
  27   # Get the file's description (all the text before the first PEM block)
  28   file_description = GetTextUntilNextPemBlock(file_data)
  29
  30   result += file_description + '\n'
  31
  32   for block in GetPemBlocks(file_data):
  33     result += '\n\n\n'
  34
  35     result += MakePemBlockString(block.name, block.data)
  36
  37     # If there was a user comment (non-script-generated comment) associated
  38     # with the block, output it immediately after the block.
  39     user_comment = GetUserComment(block.comment)
  40     if user_comment:
  41       result += '\n' + user_comment + '\n'
  42
  43     # For every block except for DATA, try to pretty print the parsed ASN.1.
  44     # DATA blocks likely would be DER in practice, but for the purposes of
  45     # these tests seeing its structure doesn't clarify
  46     # anything and is just a distraction.
  47     if block.name != 'DATA':
  48       generated_comment = GenerateCommentForBlock(block.name, block.data)
  49       result += '\n' + generated_comment + '\n'
  50
  51   return result
  52
  53
  54 def GenerateCommentForBlock(block_name, block_data):
  55   """Returns a string describing the ASN.1 structure of block_data"""
  56
  57   p = subprocess.Popen(['openssl', 'asn1parse', '-i', '-inform', 'DER'],
  58                        stdout=subprocess.PIPE, stdin=subprocess.PIPE,
  59                        stderr=subprocess.PIPE)
  60   stdout_data, stderr_data = p.communicate(input=block_data)
  61   generated_comment = '$ openssl asn1parse -i < [%s]\n%s' % (block_name,
  62                                                              stdout_data)
  63   return generated_comment.strip('\n')
  64
  65
  66 def GetTextUntilNextPemBlock(text):
  67   return text.split('-----BEGIN ', 1)[0].strip('\n')
  68
  69
  70 def GetUserComment(comment):
  71   """Removes any script-generated lines (everything after the $ openssl line)"""
  72
  73   # Consider everything after "$ openssl" to be a generated comment.
  74   comment = comment.split('$ openssl asn1parse -i', 1)[0].strip('\n')
  75   if IsEntirelyWhiteSpace(comment):
  76     comment = ''
  77   return comment
  78
  79
  80 def MakePemBlockString(name, data):
  81   return ('-----BEGIN %s-----\n'
  82           '%s'
  83           '-----END %s-----\n') % (name, EncodeDataForPem(data), name)
  84
  85
  86 def GetPemFilePaths():
  87   """Returns an iterable for all the paths to the PEM test files"""
  88
  89   base_dir = os.path.dirname(os.path.realpath(__file__))
  90   return glob.iglob(os.path.join(base_dir, '*.pem'))
  91
  92
  93 def ReadFileToString(path):
  94   with open(path, 'r') as f:
  95     return f.read()
  96
  97
  98 def WrapTextToLineWidth(text, column_width):
  99   result = ''
 100   pos = 0
 101   while pos < len(text):
 102     result += text[pos : pos + column_width] + '\n'
 103     pos += column_width
 104   return result
 105
 106
 107 def EncodeDataForPem(data):
 108   result = base64.b64encode(data)
 109   return WrapTextToLineWidth(result, 75)
 110
 111
 112 class PemBlock(object):
 113   def __init__(self):
 114     self.name = None
 115     self.data = None
 116     self.comment = None
 117
 118
 119 def StripAllWhitespace(text):
 120   pattern = re.compile(r'\s+')
 121   return re.sub(pattern, '', text)
 122
 123
 124 def IsEntirelyWhiteSpace(text):
 125   return len(StripAllWhitespace(text)) == 0
 126
 127
 128 def DecodePemBlockData(text):
 129   text = StripAllWhitespace(text)
 130   return base64.b64decode(text)
 131
 132
 133 def GetPemBlocks(data):
 134   """Returns an iterable of PemBlock"""
 135
 136   regex = re.compile(r'-----BEGIN ([\w ]+)-----(.*?)-----END \1-----',
 137                      re.DOTALL)
 138
 139   for match in regex.finditer(data):
 140     block = PemBlock()
 141
 142     block.name = match.group(1)
 143     block.data = DecodePemBlockData(match.group(2))
 144
 145     # Keep track of any non-PEM text between blocks
 146     block.comment = GetTextUntilNextPemBlock(data[match.end():])
 147
 148     yield block
 149
 150
 151 def WriteStringToFile(data, path):
 152   with open(path, "w") as f:
 153     f.write(data)
 154
 155
 156 def main():
 157   for path in GetPemFilePaths():
 158     print "Processing %s ..." % (path)
 159     original_data = ReadFileToString(path)
 160     transformed_data = Transform(original_data)
 161     if original_data != transformed_data:
 162       WriteStringToFile(transformed_data, path)
 163       print "Rewrote %s" % (path)
 164
 165
 166 if __name__ == "__main__":
 167   main()