From a49872f80882fcd29df331eeec81bfb8fa298531 Mon Sep 17 00:00:00 2001
From: Ian Gulliver <flamingcow@gmail.com>
Date: Tue, 5 Apr 2016 17:39:16 -0700
Subject: [PATCH] Add fetcher and imager support for HTTPS client/server auth.

---
 client/fetcher.py | 40 ++++++++++++++++++++++++++++++----------
 imager/image.py   | 42 ++++++++++++++++++++++++++++++++++++------
 2 files changed, 66 insertions(+), 16 deletions(-)

diff --git a/client/fetcher.py b/client/fetcher.py
index d22bb9f..3d299f7 100755
--- a/client/fetcher.py
+++ b/client/fetcher.py
@@ -6,12 +6,12 @@ import json
 import hashlib
 import os
 import re
+import requests
 import shutil
 import socket
 import struct
 import subprocess
 import tempfile
-import urllib.request
 from OpenSSL import crypto
 
 
@@ -26,6 +26,18 @@ parser.add_argument(
     dest='ca_cert',
     action='store',
     required=True)
+parser.add_argument(
+    '--https-ca-cert',
+    dest='https_ca_cert',
+    action='store')
+parser.add_argument(
+    '--https-client-cert',
+    dest='https_client_cert',
+    action='store')
+parser.add_argument(
+    '--https-client-key',
+    dest='https_client_key',
+    action='store')
 parser.add_argument(
     '--image-dir',
     dest='image_dir',
@@ -62,10 +74,15 @@ class Fetcher(object):
   _MAX_BP = 10000
   _FILE_REGEX = re.compile('^(?P<timestamp>\d+)\.iso$')
 
-  def __init__(self, base_url, ca_cert, image_dir):
+  def __init__(self, base_url, ca_cert, image_dir, https_ca_cert, https_client_cert, https_client_key):
     self._base_url = base_url
     self._ca_cert_path = ca_cert
     self._image_dir = image_dir
+    self._session = requests.Session()
+    if https_ca_cert:
+      self._session.verify = https_ca_cert
+    if https_client_cert and https_client_key:
+      self._session.cert = (https_client_cert, https_client_key)
 
   def _VerifyChain(self, untrusted_certs, cert):
     tempdir = tempfile.mkdtemp()
@@ -105,8 +122,8 @@ class Fetcher(object):
 
   def _GetManifest(self):
     url = '%s/manifest.json' % (self._base_url)
-    resp = urllib.request.urlopen(url).read().decode('utf8')
-    unwrapped = self._Unwrap(json.loads(resp))
+    resp = self._session.get(url)
+    unwrapped = self._Unwrap(resp.json())
     self._ValidateManifest(unwrapped)
     return unwrapped
 
@@ -148,15 +165,12 @@ class Fetcher(object):
 
     url = '%s/%s' % (self._base_url, filename)
     print('Fetching:', url)
-    resp = urllib.request.urlopen(url)
+    resp = self._session.get(url, stream=True)
 
     hash_obj = hashlib.sha256()
     try:
       fh = tempfile.NamedTemporaryFile(dir=self._image_dir, delete=False)
-      while True:
-        data = resp.read(self._BUF_SIZE)
-        if not data:
-          break
+      for data in resp.iter_content(self._BUF_SIZE):
         hash_obj.update(data)
         fh.write(data)
       if hash_obj.hexdigest() != image['hash']:
@@ -207,7 +221,13 @@ class Fetcher(object):
 
 
 def main():
-  fetcher = Fetcher(FLAGS.base_url, FLAGS.ca_cert, FLAGS.image_dir)
+  fetcher = Fetcher(
+      FLAGS.base_url,
+      FLAGS.ca_cert,
+      FLAGS.image_dir,
+      FLAGS.https_ca_cert,
+      FLAGS.https_client_cert,
+      FLAGS.https_client_key)
   fetcher.Fetch()
   fetcher.DeleteOldImages(FLAGS.max_images)
 
diff --git a/imager/image.py b/imager/image.py
index 19fd496..ea75fc0 100755
--- a/imager/image.py
+++ b/imager/image.py
@@ -20,6 +20,18 @@ parser.add_argument(
     dest='ca_cert',
     action='store',
     required=True)
+parser.add_argument(
+    '--https-ca-cert',
+    dest='https_ca_cert',
+    action='store')
+parser.add_argument(
+    '--https-client-cert',
+    dest='https_client_cert',
+    action='store')
+parser.add_argument(
+    '--https-client-key',
+    dest='https_client_key',
+    action='store')
 parser.add_argument(
     '--device',
     dest='device',
@@ -36,11 +48,23 @@ FLAGS = parser.parse_args()
 
 class Imager(object):
 
-  def __init__(self, device, persistent_percent, base_url, ca_cert):
+  def __init__(self, device, persistent_percent, base_url, ca_cert, https_ca_cert, https_client_cert, https_client_key):
     self._device = device
     self._persistent_percent = persistent_percent
-    self._base_url = base_url
-    self._ca_cert = ca_cert
+
+    self._fetcher_args = [
+        '--base-url', base_url,
+        '--ca-cert', ca_cert,
+    ]
+    if https_ca_cert:
+      self._fetcher_args.extend([
+          '--https-ca-cert', https_ca_cert,
+      ])
+    if https_client_cert and https_client_key:
+      self._fetcher_args.extend([
+          '--https-client-cert', https_client_cert,
+          '--https-client-key', https_client_key,
+      ])
 
     self._icon_path = os.path.dirname(sys.argv[0])
 
@@ -123,8 +147,7 @@ class Imager(object):
     self._Exec(
         fetcher,
         '--image-dir', image_path,
-        '--base-url', self._base_url,
-        '--ca-cert', self._ca_cert)
+        *self._fetcher_args)
 
     return image_path
 
@@ -160,7 +183,14 @@ class Imager(object):
 
 
 def main():
-  imager = Imager(FLAGS.device, FLAGS.persistent_percent, FLAGS.base_url, FLAGS.ca_cert)
+  imager = Imager(
+      FLAGS.device,
+      FLAGS.persistent_percent,
+      FLAGS.base_url,
+      FLAGS.ca_cert,
+      FLAGS.https_ca_cert,
+      FLAGS.https_client_cert,
+      FLAGS.https_client_key)
   imager.Image()