Add downloader function for trezors

Author: achow101

hwilib/firmware.py

@@ -1,12 +1,63 @@
-#! /usr/bin/env bash
+# Firmware download things
 
 import json
 import logging
+import os
+import requests
 import sys
 
+from urllib.parse import urlparse
+
 from . import __version__
 from .cli import HWIArgumentParser
-from .errors import handle_errors
+from .errors import BadArgumentError, handle_errors, UnknownDeviceError
+
+def _download_file(url):
+    filename = os.path.basename(urlparse(url).path)
+
+    with requests.get(url, stream=True) as r:
+        r.raise_for_status()
+        with open(filename, 'wb') as f:
+            for chunk in r.iter_content(chunk_size=8192):
+                if chunk:
+                    f.write(chunk)
+
+    return os.path.abspath(filename)
+
+def _trezor_download(version=None, bitcoinonly=False, device_version=1):
+    releases = requests.get('https://wallet.trezor.io/data/firmware/{}/releases.json'.format(device_version)).json()
+    if not releases:
+        raise IOError('Could not get list of releases')
+
+    if bitcoinonly:
+        releases = [r for r in releases if "url_bitcoinonly" in r]
+    releases.sort(key=lambda r: r["version"], reverse=True)
+
+    version_info = {}
+    if version is None:
+        version_info = releases[0]
+        version = '.'.join([str(x) for x in version_info['version']])
+    else:
+        version_list = [int(x) for x in version.split(".")]
+        for r in releases:
+            if r['version'] == version_list:
+                version_info = r
+                break
+        else:
+            raise BadArgumentError('{} is not available'.format(version))
+
+    url = 'https://wallet.trezor.io/{}'.format(version_info['url_bitcoinonly'] if bitcoinonly else version_info['url'])
+    downloaded_file = _download_file(url)
+
+    bitcoinonly_str = 'Bitcoin only ' if bitcoinonly else ''
+
+    return {'success': True, 'message': 'Trezor 1 {}firmware version {} downloaded to {}'.format(bitcoinonly_str, version, downloaded_file)}
+
+def trezor_1_download(version=None, bitcoinonly=False):
+    return _trezor_download(version, bitcoinonly, 1)
+
+def trezor_t_download(version=None, bitcoinonly=False):
+    return _trezor_download(version, bitcoinonly, 2)
 
 def download_firmware(model, version, bitcoinonly=False):
     dev_model = model.lower()