Validate origins when generating subdomain tokens

tools/origin_trials/generate_token.py

 #!/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.
 
 """Utility for generating experimental API tokens
 
 usage: generate_token.py [-h] [--key-file KEY_FILE]
                          [--expire-days EXPIRE_DAYS |
                           --expire-timestamp EXPIRE_TIMESTAMP]
                          [--is_subdomain | --no-subdomain]
                          origin trial_name
 
 Run "generate_token.py -h" for more help on usage.
 """
 import argparse
 import base64
 from datetime import datetime
 import json
 import re
 import os
 import struct
+import subprocess
 import sys
 import time
 import urlparse
 
 script_dir = os.path.dirname(os.path.realpath(__file__))
 sys.path.insert(0, os.path.join(script_dir, 'third_party', 'ed25519'))
 import ed25519
 
 
 # Matches a valid DNS name label (alphanumeric plus hyphens, except at the ends,
 # no longer than 63 ASCII characters)
 DNS_LABEL_REGEX = re.compile(r"^(?!-)[a-z\d-]{1,63}(?<!-)$", re.IGNORECASE)
 
 # This script generates Version 2 tokens.
 VERSION = "\x02"
 
 # Default key file, relative to script_dir.
 DEFAULT_KEY_FILE = 'eftest.key'
 
+# Default location of validate subdomain utility, relative to script_dir.
+DEFAULT_TARGET_PATH = '../../out/Default/'
+
 def HostnameFromArg(arg):
   """Determines whether a string represents a valid hostname.
 
   Returns the canonical hostname if its argument is valid, or None otherwise.
   """
   if not arg or len(arg) > 255:
     return None
   if arg[-1] == ".":
     arg = arg[:-1]
   if "." not in arg and arg != "localhost":
@@ skipping 27 matching lines @@
   if not port:
     port = {"https": 443, "http": 80}[origin.scheme]
   # Strip any extra components and return the origin URL:
   return "{0}://{1}:{2}".format(origin.scheme, origin.hostname, port)
 
 def ExpiryFromArgs(args):
   if args.expire_timestamp:
     return int(args.expire_timestamp)
   return (int(time.time()) + (int(args.expire_days) * 86400))
 
+def ValidateSubdomainTokenOrigin(origin, target_path):
+  """ Calls validate_subdomain_origin utility to check the origin
+
+  If the utility is not found, prints a warning for manual validation, and
+  returns True
+  """
+  utility_path = "%s/validate_subdomain_origin" % target_path
+  if not os.path.exists(utility_path):
+    print "WARNING!"
+    print "Origin not validated for use in subdomain token"
+    print "  (missing '%s' utility)" % utility_path
+    print "Must manually check origin against the Public Suffix List"
+    print
+    return True
+
+  rc = subprocess.call([utility_path, "--quiet", origin])
+  if (rc < 0 or rc > 4):
+    print("Unexpected return code from %s: %d" % (utility_path, rc))
+    sys.exit(1)
+
+  return rc == 0
+
 def GenerateTokenData(origin, is_subdomain, feature_name, expiry):
   data = {"origin": origin,
           "feature": feature_name,
           "expiry": expiry}
   if is_subdomain is not None:
     data["isSubdomain"] = is_subdomain
   return json.dumps(data).encode('utf-8')
 
 def GenerateDataToSign(version, data):
   return version + struct.pack(">I",len(data)) + data
 
 def Sign(private_key, data):
   return ed25519.signature(data, private_key[:32], private_key[32:])
 
 def FormatToken(version, signature, data):
   return base64.b64encode(version + signature +
                           struct.pack(">I",len(data)) + data)
 
 def main():
   default_key_file_absolute = os.path.join(script_dir, DEFAULT_KEY_FILE)
+  default_target_path_absolute = os.path.join(script_dir, DEFAULT_TARGET_PATH)
 
   parser = argparse.ArgumentParser(
       description="Generate tokens for enabling experimental features")
   parser.add_argument("origin",
                       help="Origin for which to enable the feature. This can "
                            "be either a hostname (default scheme HTTPS, "
                            "default port 443) or a URL.",
                       type=OriginFromArg)
   parser.add_argument("trial_name",
                       help="Feature to enable. The current list of "
@@ skipping 19 matching lines @@
   expiry_group = parser.add_mutually_exclusive_group()
   expiry_group.add_argument("--expire-days",
                             help="Days from now when the token should expire",
                             type=int,
                             default=42)
   expiry_group.add_argument("--expire-timestamp",
                             help="Exact time (seconds since 1970-01-01 "
                                  "00:00:00 UTC) when the token should expire",
                             type=int)
 
+  parser.add_argument("--target",
+                      help="Path to the output directory for compiled "
+                           "resources",
+                      default=default_target_path_absolute)
+
   args = parser.parse_args()
   expiry = ExpiryFromArgs(args)
 
   key_file = open(os.path.expanduser(args.key_file), mode="rb")
   private_key = key_file.read(64)
 
   # Validate that the key file read was a proper Ed25519 key -- running the
   # publickey method on the first half of the key should return the second
   # half.
   if (len(private_key) < 64 or
     ed25519.publickey(private_key[:32]) != private_key[32:]):
     print("Unable to use the specified private key file.")
     sys.exit(1)
 
+  # For subdomain tokens, validate that the origin is allowed
+  if args.is_subdomain:
+    target_path = os.path.expanduser(args.target)
+    if not ValidateSubdomainTokenOrigin(args.origin, target_path):
+      print "The specified origin is not valid for use in a subdomain token."
+      sys.exit(1)
+
   token_data = GenerateTokenData(args.origin, args.is_subdomain,
                                  args.trial_name, expiry)
   data_to_sign = GenerateDataToSign(VERSION, token_data)
   signature = Sign(private_key, data_to_sign)
 
   # Verify that that the signature is correct before printing it.
   try:
     ed25519.checkvalid(signature, data_to_sign, private_key[32:])
   except Exception, exc:
     print "There was an error generating the signature."
     print "(The original error was: %s)" % exc
     sys.exit(1)
 
 
   # Output the token details
   print "Token details:"
   print " Origin: %s" % args.origin
   print " Is Subdomain: %s" % args.is_subdomain
   print " Feature: %s" % args.trial_name
   print " Expiry: %d (%s UTC)" % (expiry, datetime.utcfromtimestamp(expiry))
   print
 
   # Output the properly-formatted token.
   print FormatToken(VERSION, signature, token_data)
 
 if __name__ == "__main__":
   main()