Replace aiohttp with httpx and add the next steps for update_firmware

This commit is contained in:
1e9abhi1e10
2024-05-26 02:22:21 +05:30
parent 1ab05c7a5e
commit 2d92c2c0e2

View File

@@ -16,7 +16,7 @@
import asyncio import asyncio
import ipaddress import ipaddress
import warnings import warnings
import aiohttp import httpx
import tempfile import tempfile
import os import os
from typing import List, Optional, Protocol, Tuple, Type, TypeVar, Union from typing import List, Optional, Protocol, Tuple, Type, TypeVar, Union
@@ -548,16 +548,23 @@ class MinerProtocol(Protocol):
if current_firmware_version == latest_firmware_info['version']: if current_firmware_version == latest_firmware_info['version']:
return True # Firmware is already up to date return True # Firmware is already up to date
# Download the new firmware file # Download the new firmware file if it's not a local file
if not firmware_url.startswith("file://"):
firmware_file_path = await self._download_firmware(latest_firmware_info['download_url']) firmware_file_path = await self._download_firmware(latest_firmware_info['download_url'])
if not firmware_file_path: if not firmware_file_path:
raise IOError("Failed to download the firmware file.") raise IOError("Failed to download the firmware file.")
else:
firmware_file_path = latest_firmware_info['download_url']
# Transfer the firmware file to the miner # Transfer the firmware file to the miner
# TODO transfer_success = await self._transfer_firmware_to_miner(firmware_file_path)
if not transfer_success:
raise IOError("Failed to transfer the firmware file to the miner.")
# Apply the firmware update on the miner # Apply the firmware update on the miner
# TODO update_success = await self._apply_firmware_update(firmware_file_path)
if not update_success:
raise IOError("Failed to apply the firmware update.")
# Reboot the miner # Reboot the miner
# TODO # TODO
@@ -596,11 +603,11 @@ class BaseMiner(MinerProtocol):
Returns: Returns:
A dictionary containing the firmware version and download URL. A dictionary containing the firmware version and download URL.
""" """
async with aiohttp.ClientSession() as session: async with httpx.AsyncClient() as client:
async with session.get(firmware_url) as response: response = await client.get(firmware_url)
if response.status != 200: if response.status_code != 200:
raise ConnectionError(f"Failed to fetch firmware info, status code: {response.status}") raise ConnectionError(f"Failed to fetch firmware info, status code: {response.status_code}")
firmware_info = await response.json() firmware_info = response.json()
return { return {
'version': firmware_info['version'], 'version': firmware_info['version'],
'download_url': firmware_info['download_url'] 'download_url': firmware_info['download_url']
@@ -616,19 +623,23 @@ class BaseMiner(MinerProtocol):
Returns: Returns:
The file path to the downloaded firmware file. The file path to the downloaded firmware file.
""" """
async with aiohttp.ClientSession() as session: async with httpx.AsyncClient() as client:
async with session.get(download_url) as response: response = await client.get(download_url)
if response.status != 200: if response.status_code != 200:
raise ConnectionError(f"Failed to download firmware, status code: {response.status}") raise ConnectionError(f"Failed to download firmware, status code: {response.status_code}")
# Assuming the Content-Disposition header contains the filename # Assuming the Content-Disposition header contains the filename
content_disposition = response.headers.get('Content-Disposition') content_disposition = response.headers.get('Content-Disposition')
if not content_disposition: if not content_disposition:
raise ValueError("Could not determine the filename of the firmware download.") raise ValueError("Could not determine the filename of the firmware download.")
filename = content_disposition.split("filename=")[-1].strip().strip('"') filename = content_disposition.split("filename=")[-1].strip().strip('"')
temp_dir = tempfile.gettempdir() temp_dir = tempfile.gettempdir()
file_path = os.path.join(temp_dir, filename) file_path = os.path.join(temp_dir, filename)
with open(file_path, 'wb') as file: with open(file_path, 'wb') as file:
file.write(await response.read()) file.write(response.content)
return file_path return file_path
async def _query_firmware_server(self, server_url: str) -> dict: async def _query_firmware_server(self, server_url: str) -> dict:
@@ -641,11 +652,11 @@ class BaseMiner(MinerProtocol):
Returns: Returns:
A dictionary containing the firmware version and download URL. A dictionary containing the firmware version and download URL.
""" """
async with aiohttp.ClientSession() as session: async with httpx.AsyncClient() as client:
async with session.get(server_url) as response: response = await client.get(server_url)
if response.status != 200: if response.status_code != 200:
raise ConnectionError(f"Failed to query firmware server, status code: {response.status}") raise ConnectionError(f"Failed to query firmware server, status code: {response.status_code}")
firmware_files = await response.json() firmware_files = response.json()
for firmware_file in firmware_files: for firmware_file in firmware_files:
if self._is_appropriate_firmware(firmware_file['version']): if self._is_appropriate_firmware(firmware_file['version']):
return { return {
@@ -654,5 +665,93 @@ class BaseMiner(MinerProtocol):
} }
raise ValueError("No appropriate firmware file found on the server.") raise ValueError("No appropriate firmware file found on the server.")
async def _transfer_firmware_to_miner(self, firmware_file_path: str) -> bool:
"""
Transfer the firmware file to the miner using SSH.
Parameters:
firmware_file_path: The file path to the firmware file that needs to be transferred.
Returns:
A boolean value indicating the success of the transfer.
"""
import paramiko
from paramiko import SSHClient
from scp import SCPClient
# Retrieve SSH credentials from environment variables or secure storage
ssh_username = os.getenv('MINER_SSH_USERNAME')
ssh_password = os.getenv('MINER_SSH_PASSWORD')
ssh_key_path = os.getenv('MINER_SSH_KEY_PATH')
try:
# Establish an SSH client session using credentials
with SSHClient() as ssh_client:
ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
if ssh_key_path:
ssh_client.connect(self.ip, username=ssh_username, key_filename=ssh_key_path)
else:
ssh_client.connect(self.ip, username=ssh_username, password=ssh_password)
# Use SCPClient to transfer the firmware file to the miner
with SCPClient(ssh_client.get_transport()) as scp_client:
scp_client.put(firmware_file_path, remote_path='/tmp/')
return True
except Exception as e:
logger.error(f"Failed to transfer firmware to miner {self.ip}: {e}")
return False
async def _apply_firmware_update(self, firmware_file_path: str) -> bool:
"""
Apply the firmware update to the miner.
Parameterrs:
firmware_file_path: The path to the firmware file to be uploaded.
Returns:
True if the firmware update was successful, False otherwise.
"""
try:
# Identify the miner using the get_miner function
miner = await self.get_miner(self.ip)
if not miner:
logger.error(f"Failed to identify miner at {self.ip}")
return False
# Gather necessary data from the miner
miner_data = await miner.get_data()
# Logic to apply firmware update using SSH
if miner.ssh:
async with miner.ssh as ssh_client:
await ssh_client.upload_file(firmware_file_path, "/tmp/firmware.bin")
await ssh_client.execute_command("upgrade_firmware /tmp/firmware.bin")
return True
# Logic to apply firmware update using RPC
elif miner.rpc:
await miner.rpc.upload_firmware(firmware_file_path)
await miner.rpc.execute_command("upgrade_firmware")
return True
else:
logger.error("No valid interface available to apply firmware update.")
return False
except Exception as e:
logger.error(f"Error applying firmware update: {e}")
return False
async def _is_appropriate_firmware(self, firmware_version: str) -> bool:
"""
Determine if the given firmware version is appropriate for the miner.
Parameters:
firmware_version: The version of the firmware to evaluate.
Returns:
A boolean indicating if the firmware version is appropriate.
"""
# TODO
return True
AnyMiner = TypeVar("AnyMiner", bound=BaseMiner) AnyMiner = TypeVar("AnyMiner", bound=BaseMiner)