ipowerswitch/ipowerswitch.py

import dataclasses
import enum
import logging
import re
import typing

import bs4
import requests
import requests.auth
import rich.console
import rich.table
import typer
import yarl


class SwitchStatus(str, enum.Enum):
    on = "on"
    off = "off"


@dataclasses.dataclass
class Switch:
    address: str
    user: str
    password: str

    def switch(
        self,
        switch_id: int,
        outlet: int,
        from_state: SwitchStatus | None = None,
    ) -> None:
        url = yarl.URL(f"http://{self.address}").joinpath(
            "cgi-bin",
            f"iswitch{switch_id:02d}",
            "irswitch.exe",
        )

        params = {
            "CURRENT": f"{switch_id:02d}",
            f"SW{outlet}.x": 1,
            f"SW{outlet}.y": 1,
        }
        if from_state:
            params[f"STATUS{outlet}"] = f"{from_state.value.upper():3s}"

        logging.info("Switching outlet: %s", params)
        response = requests.post(
            str(url),
            auth=requests.auth.HTTPBasicAuth(self.user, self.password),
            data=params,
            allow_redirects=False,
        )
        print(response)

    def switch_on(self, switch_id: int, outlet: int) -> None:
        return self.switch(switch_id, outlet, from_state=SwitchStatus.off)

    def switch_off(self, switch_id: int, outlet: int) -> None:
        return self.switch(switch_id, outlet, from_state=SwitchStatus.on)

    def switch_status(self, switch_id: int) -> dict[int, SwitchStatus]:
        url = yarl.URL(f"http://{self.address}").joinpath(
            f"iswitch{switch_id:02d}",
            "index.htm",
        )
        response = requests.get(
            str(url),
            auth=requests.auth.HTTPBasicAuth(self.user, self.password),
        )
        soup = bs4.BeautifulSoup(response.content, features="html.parser")
        inputs = {i["name"]: i["value"] for i in soup.find_all("input", type="hidden")}
        outlets = int(inputs["OUTLET"])
        return {
            outlet: SwitchStatus[inputs.get(f"STATUS{outlet}", "off").strip().lower()]
            for outlet in range(1, outlets + 1)
        }


app = typer.Typer()
console = rich.console.Console()

SwitchIdArgument = typing.Annotated[int, typer.Argument(min=1, max=16)]
OutletArgument = typing.Annotated[int, typer.Argument(min=1, max=8)]


@app.command()
def switch(
    context: typer.Context,
    switch_id: SwitchIdArgument,
    outlet: OutletArgument,
    action: typing.Annotated[SwitchStatus, typer.Argument(case_sensitive=False)],
) -> None:
    switch: Switch = context.obj
    match action:
        case SwitchStatus.on:
            switch.switch_on(switch_id, outlet)
        case SwitchStatus.off:
            switch.switch_off(switch_id, outlet)


@app.command()
def status(context: typer.Context, switch_id: SwitchIdArgument) -> None:
    switch: Switch = context.obj
    statuses = switch.switch_status(switch_id)
    if not console.is_interactive:
        for key in sorted(statuses.keys()):
            console.print(f"{key}\t{statuses[key].value.upper()}")
    else:
        table = rich.table.Table("Outlet", "Status")
        for key in sorted(statuses.keys()):
            table.add_row(str(key), statuses[key].value.upper())
            console.print(table)


@app.callback()
def main(
    context: typer.Context,
    address: str,
    user: str = "admin",
    password: str = "admin",
) -> None:
    logging.basicConfig(level=logging.INFO)
    context.obj = Switch(
        address=address,
        user=user,
        password=password,
    )


if __name__ == "__main__":
    app()