verify checksums

teeth_agent/standby.py

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import hashlib
 import json
 import os
 import subprocess
@@ -81,8 +82,15 @@ def _download_image(image_info):
 
 
 def _verify_image(image_info, image_location):
-    # TODO(jimrollenhagen) verify the image checksum
-    return True
+    hashes = image_info['hashes']
+    for k, v in hashes.iteritems():
+        algo = getattr(hashlib, k, None)
+        if algo is None:
+            continue
+        hash_ = algo(open(image_location).read())
+        if hash_ == v:
+            return True
+    return False
 
 
 class CacheImagesCommand(base.AsyncCommandResult):

teeth_agent/tests/standby_agent.py

@@ -121,15 +121,19 @@ class TestBaseTeethAgent(unittest.TestCase):
 
         call_mock.assert_called_once_with(command)
 
+    @mock.patch('hashlib.md5', autospec=True)
     @mock.patch('__builtin__.open', autospec=True)
     @mock.patch('requests.get', autospec=True)
-    def test_download_image(self, requests_mock, open_mock):
+    def test_download_image(self, requests_mock, open_mock, md5_mock):
         image_info = self._build_fake_image_info()
         response = requests_mock.return_value
         response.status_code = 200
         response.iter_content.return_value = ['some', 'content']
         open_mock.return_value.__enter__ = lambda s: s
         open_mock.return_value.__exit__ = mock.Mock()
+        read_mock = open_mock.return_value.read
+        read_mock.return_value = 'content'
+        md5_mock.return_value = image_info['hashes'].values()[0]
 
         standby._download_image(image_info)
         requests_mock.assert_called_once_with(image_info['urls'][0],