pyasic

class MinerFactory:
    def __init__(self):
        self.cache = {}

    def clear_cached_miners(self):
        self.cache = {}

    async def get_multiple_miners(self, ips: List[str], limit: int = 200):
        results = []

        async for miner in self.get_miner_generator(ips, limit):
            results.append(miner)

        return results

    async def get_miner_generator(self, ips: list, limit: int = 200):
        tasks = []
        semaphore = asyncio.Semaphore(limit)

        for ip in ips:
            tasks.append(asyncio.create_task(self.get_miner(ip)))

        for task in tasks:
            await semaphore.acquire()
            try:
                result = await task
                if result is not None:
                    yield result
            finally:
                semaphore.release()

    async def get_miner(self, ip: str):
        ip = str(ip)
        if ip in self.cache:
            return self.cache[ip]

        miner_type = None

        for _ in range(RETRIES):
            task = asyncio.create_task(self._get_miner_type(ip))
            try:
                miner_type = await asyncio.wait_for(task, timeout=TIMEOUT)
            except asyncio.TimeoutError:
                task.cancel()
            else:
                if miner_type is not None:
                    break

        if miner_type is not None:
            miner_model = None
            miner_model_fns = {
                MinerTypes.ANTMINER: self.get_miner_model_antminer,
                MinerTypes.WHATSMINER: self.get_miner_model_whatsminer,
                MinerTypes.AVALONMINER: self.get_miner_model_avalonminer,
                MinerTypes.INNOSILICON: self.get_miner_model_innosilicon,
                MinerTypes.GOLDSHELL: self.get_miner_model_goldshell,
                MinerTypes.BRAIINS_OS: self.get_miner_model_braiins_os,
                MinerTypes.VNISH: self.get_miner_model_vnish,
                MinerTypes.HIVEON: self.get_miner_model_hiveon,
                MinerTypes.LUX_OS: self.get_miner_model_luxos,
            }
            fn = miner_model_fns.get(miner_type)

            if fn is not None:
                task = asyncio.create_task(fn(ip))
                try:
                    miner_model = await asyncio.wait_for(task, timeout=30)
                except asyncio.TimeoutError:
                    task.cancel()

            miner = self._select_miner_from_classes(
                ip, miner_type=miner_type, miner_model=miner_model
            )

            if miner is not None and not isinstance(miner, UnknownMiner):
                self.cache[ip] = miner
            return miner

    async def _get_miner_type(self, ip: str):
        tasks = [
            asyncio.create_task(self._get_miner_web(ip)),
            asyncio.create_task(self._get_miner_socket(ip)),
        ]

        return await concurrent_get_first_result(tasks, lambda x: x is not None)

    async def _get_miner_web(self, ip: str):
        urls = [f"http://{ip}/", f"https://{ip}/"]
        async with httpx.AsyncClient(verify=False) as session:
            tasks = [asyncio.create_task(self._web_ping(session, url)) for url in urls]

            text, resp = await concurrent_get_first_result(
                tasks, lambda x: x[0] is not None
            )
            if text is not None:
                return self._parse_web_type(text, resp)

    @staticmethod
    async def _web_ping(
        session: httpx.AsyncClient, url: str
    ) -> Tuple[Optional[str], Optional[httpx.Response]]:
        try:
            resp = await session.get(url, follow_redirects=False)
            return resp.text, resp
        except (httpx.HTTPError, asyncio.TimeoutError):
            pass
        return None, None

    @staticmethod
    def _parse_web_type(web_text: str, web_resp: httpx.Response) -> MinerTypes:
        if web_resp.status_code == 401 and 'realm="antMiner' in web_resp.headers.get(
            "www-authenticate", ""
        ):
            return MinerTypes.ANTMINER
        if web_resp.status_code == 307 and "https://" in web_resp.headers.get(
            "location", ""
        ):
            return MinerTypes.WHATSMINER
        if "Braiins OS" in web_text:
            return MinerTypes.BRAIINS_OS
        if "cloud-box" in web_text:
            return MinerTypes.GOLDSHELL
        if "AnthillOS" in web_text:
            return MinerTypes.VNISH
        if "Avalon" in web_text:
            return MinerTypes.AVALONMINER
        if "DragonMint" in web_text:
            return MinerTypes.INNOSILICON

    async def _get_miner_socket(self, ip: str):
        commands = ["version", "devdetails"]
        tasks = [asyncio.create_task(self._socket_ping(ip, cmd)) for cmd in commands]

        data = await concurrent_get_first_result(
            tasks, lambda x: x is not None and self._parse_socket_type(x) is not None
        )
        if data is not None:
            d = self._parse_socket_type(data)
            return d

    @staticmethod
    async def _socket_ping(ip: str, cmd: str) -> Optional[str]:
        data = b""
        try:
            reader, writer = await asyncio.wait_for(
                asyncio.open_connection(str(ip), 4028), timeout=30
            )
        except (ConnectionError, OSError, asyncio.TimeoutError):
            return

        cmd = {"command": cmd}

        try:
            # send the command
            writer.write(json.dumps(cmd).encode("utf-8"))
            await writer.drain()

            # loop to receive all the data
            while True:
                try:
                    d = await asyncio.wait_for(reader.read(4096), timeout=1)
                    if not d:
                        break
                    data += d
                except asyncio.TimeoutError:
                    pass
                except ConnectionResetError:
                    return
        except asyncio.CancelledError:
            raise
        except (ConnectionError, OSError):
            return
        finally:
            # Handle cancellation explicitly
            if writer.transport.is_closing():
                writer.transport.close()
            else:
                writer.close()
            try:
                await writer.wait_closed()
            except (ConnectionError, OSError):
                return
        if data:
            return data.decode("utf-8")

    @staticmethod
    def _parse_socket_type(data: str) -> MinerTypes:
        upper_data = data.upper()
        if "BOSMINER" in upper_data or "BOSER" in upper_data:
            return MinerTypes.BRAIINS_OS
        if "BTMINER" in upper_data or "BITMICRO" in upper_data:
            return MinerTypes.WHATSMINER
        if "VNISH" in upper_data:
            return MinerTypes.VNISH
        if "HIVEON" in upper_data:
            return MinerTypes.HIVEON
        if "LUXMINER" in upper_data:
            return MinerTypes.LUX_OS
        if "ANTMINER" in upper_data and not "DEVDETAILS" in upper_data:
            return MinerTypes.ANTMINER
        if "INTCHAINS_QOMO" in upper_data:
            return MinerTypes.GOLDSHELL
        if "AVALON" in upper_data:
            return MinerTypes.AVALONMINER

    async def send_web_command(
        self,
        ip: Union[ipaddress.ip_address, str],
        location: str,
        auth: Optional[httpx.DigestAuth] = None,
    ) -> Optional[dict]:
        async with httpx.AsyncClient(verify=False) as session:
            try:
                data = await session.get(
                    f"http://{str(ip)}{location}",
                    auth=auth,
                    timeout=30,
                )
            except (httpx.HTTPError, asyncio.TimeoutError):
                logger.info(f"{ip}: Web command timeout.")
                return
        if data is None:
            return
        try:
            json_data = data.json()
        except (json.JSONDecodeError, asyncio.TimeoutError):
            try:
                return json.loads(data.text)
            except (json.JSONDecodeError, httpx.HTTPError):
                return
        else:
            return json_data

    async def send_api_command(
        self, ip: Union[ipaddress.ip_address, str], command: str
    ) -> Optional[dict]:
        data = b""
        try:
            reader, writer = await asyncio.open_connection(str(ip), 4028)
        except (ConnectionError, OSError):
            return
        cmd = {"command": command}

        try:
            # send the command
            writer.write(json.dumps(cmd).encode("utf-8"))
            await writer.drain()

            # loop to receive all the data
            while True:
                d = await reader.read(4096)
                if not d:
                    break
                data += d

            writer.close()
            await writer.wait_closed()
        except asyncio.CancelledError:
            writer.close()
            await writer.wait_closed()
            return
        except (ConnectionError, OSError):
            return
        if data == b"Socket connect failed: Connection refused\n":
            return

        data = await self._fix_api_data(data)

        try:
            data = json.loads(data)
        except json.JSONDecodeError:
            return {}

        return data

    @staticmethod
    async def _fix_api_data(data: bytes):
        if data.endswith(b"\x00"):
            str_data = data.decode("utf-8")[:-1]
        else:
            str_data = data.decode("utf-8")
        # fix an error with a btminer return having an extra comma that breaks json.loads()
        str_data = str_data.replace(",}", "}")
        # fix an error with a btminer return having a newline that breaks json.loads()
        str_data = str_data.replace("\n", "")
        # fix an error with a bmminer return not having a specific comma that breaks json.loads()
        str_data = str_data.replace("}{", "},{")
        # fix an error with a bmminer return having a specific comma that breaks json.loads()
        str_data = str_data.replace("[,{", "[{")
        # fix an error with a btminer return having a missing comma. (2023-01-06 version)
        str_data = str_data.replace('""temp0', '","temp0')
        # fix an error with Avalonminers returning inf and nan
        str_data = str_data.replace("info", "1nfo")
        str_data = str_data.replace("inf", "0")
        str_data = str_data.replace("1nfo", "info")
        str_data = str_data.replace("nan", "0")
        # fix whatever this garbage from avalonminers is `,"id":1}`
        if str_data.startswith(","):
            str_data = f"{{{str_data[1:]}"
        # try to fix an error with overflowing the recieve buffer
        # this can happen in cases such as bugged btminers returning arbitrary length error info with 100s of errors.
        if not str_data.endswith("}"):
            str_data = ",".join(str_data.split(",")[:-1]) + "}"

        # fix a really nasty bug with whatsminer API v2.0.4 where they return a list structured like a dict
        if re.search(r"\"error_code\":\[\".+\"\]", str_data):
            str_data = str_data.replace("[", "{").replace("]", "}")

        return str_data

    @staticmethod
    def _select_miner_from_classes(
        ip: ipaddress.ip_address,
        miner_model: Union[str, None],
        miner_type: Union[MinerTypes, None],
    ) -> AnyMiner:
        try:
            return MINER_CLASSES[miner_type][str(miner_model).upper()](ip)
        except LookupError:
            if miner_type in MINER_CLASSES:
                return MINER_CLASSES[miner_type][None](ip)
            return UnknownMiner(str(ip))

    async def get_miner_model_antminer(self, ip: str):
        tasks = [
            asyncio.create_task(self._get_model_antminer_web(ip)),
            asyncio.create_task(self._get_model_antminer_sock(ip)),
        ]

        return await concurrent_get_first_result(tasks, lambda x: x is not None)

    async def _get_model_antminer_web(self, ip: str):
        # last resort, this is slow
        auth = httpx.DigestAuth("root", "root")
        web_json_data = await self.send_web_command(
            ip, "/cgi-bin/get_system_info.cgi", auth=auth
        )

        try:
            miner_model = web_json_data["minertype"]

            return miner_model
        except (TypeError, LookupError):
            pass

    async def _get_model_antminer_sock(self, ip: str):
        sock_json_data = await self.send_api_command(ip, "version")
        try:
            miner_model = sock_json_data["VERSION"][0]["Type"]

            if " (" in miner_model:
                split_miner_model = miner_model.split(" (")
                miner_model = split_miner_model[0]

            return miner_model
        except (TypeError, LookupError):
            pass

        sock_json_data = await self.send_api_command(ip, "stats")
        try:
            miner_model = sock_json_data["STATS"][0]["Type"]

            if " (" in miner_model:
                split_miner_model = miner_model.split(" (")
                miner_model = split_miner_model[0]

            return miner_model
        except (TypeError, LookupError):
            pass

    async def get_miner_model_goldshell(self, ip: str):
        json_data = await self.send_web_command(ip, "/mcb/status")

        try:
            miner_model = json_data["model"].replace("-", " ")

            return miner_model
        except (TypeError, LookupError):
            pass

    async def get_miner_model_whatsminer(self, ip: str):
        sock_json_data = await self.send_api_command(ip, "devdetails")
        try:
            miner_model = sock_json_data["DEVDETAILS"][0]["Model"]

            return miner_model
        except (TypeError, LookupError):
            pass

    async def get_miner_model_avalonminer(self, ip: str) -> Optional[str]:
        sock_json_data = await self.send_api_command(ip, "version")
        try:
            miner_model = sock_json_data["VERSION"][0]["PROD"]
            if "-" in miner_model:
                miner_model = miner_model.split("-")[0]

            return miner_model
        except (TypeError, LookupError):
            pass

    async def get_miner_model_innosilicon(self, ip: str) -> Optional[str]:
        try:
            async with httpx.AsyncClient(verify=False) as session:
                auth_req = await session.post(
                    f"http://{ip}/api/auth",
                    data={"username": "admin", "password": "admin"},
                )
                auth = auth_req.json()["jwt"]

                web_data = (
                    await session.post(
                        f"http://{ip}/api/type",
                        headers={"Authorization": "Bearer " + auth},
                        data={},
                    )
                ).json()
                return web_data["type"]
        except (httpx.HTTPError, LookupError):
            pass

    async def get_miner_model_braiins_os(self, ip: str) -> Optional[str]:
        sock_json_data = await self.send_api_command(ip, "devdetails")
        try:
            miner_model = sock_json_data["DEVDETAILS"][0]["Model"].replace(
                "Bitmain ", ""
            )

            return miner_model
        except (TypeError, LookupError):
            pass

        try:
            async with httpx.AsyncClient(verify=False) as session:
                d = await session.post(
                    f"http://{ip}/graphql",
                    json={"query": "{bosminer {info{modelName}}}"},
                )
            if d.status_code == 200:
                json_data = d.json()
                miner_model = json_data["data"]["bosminer"]["info"]["modelName"]
                return miner_model
        except (httpx.HTTPError, LookupError):
            pass

    async def get_miner_model_vnish(self, ip: str) -> Optional[str]:
        sock_json_data = await self.send_api_command(ip, "stats")
        try:
            miner_model = sock_json_data["STATS"][0]["Type"]
            if " (" in miner_model:
                split_miner_model = miner_model.split(" (")
                miner_model = split_miner_model[0]

            if "(88)" in miner_model:
                miner_model = miner_model.replace("(88)", "NOPIC")

            if " AML" in miner_model:
                miner_model = miner_model.replace(" AML", "")

            return miner_model
        except (TypeError, LookupError):
            pass

    async def get_miner_model_hiveon(self, ip: str) -> Optional[str]:
        sock_json_data = await self.send_api_command(ip, "version")
        try:
            miner_type = sock_json_data["VERSION"][0]["Type"]

            return miner_type.replace(" HIVEON", "")
        except (TypeError, LookupError):
            pass

    async def get_miner_model_luxos(self, ip: str):
        sock_json_data = await self.send_api_command(ip, "version")
        try:
            miner_model = sock_json_data["VERSION"][0]["Type"]

            if " (" in miner_model:
                split_miner_model = miner_model.split(" (")
                miner_model = split_miner_model[0]
            return miner_model
        except (TypeError, LookupError):
            pass