sporestack-python/src/sporestack/cli.py

"""
SporeStack CLI: `sporestack`
"""

import importlib.util
import json
import logging
import os
import sys
import time
from pathlib import Path
from types import ModuleType
from typing import TYPE_CHECKING, Any, Dict, Optional

if sys.version_info[:2] >= (3, 8):  # pragma: nocover
    from importlib.metadata import version as importlib_metadata_version
else:  # pragma: nocover
    # Python 3.7 doesn't have this.
    from importlib_metadata import version as importlib_metadata_version

import typer


def lazy_import(name: str) -> ModuleType:
    """
    Lazily import a module. Helps speed up CLI performance.
    """
    spec = importlib.util.find_spec(name)
    assert spec is not None
    assert spec.loader is not None
    loader = importlib.util.LazyLoader(spec.loader)
    spec.loader = loader
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    loader.exec_module(module)
    return module


# For mypy
if TYPE_CHECKING:
    from . import api_client
else:
    api_client = lazy_import("sporestack.api_client")

HELP = """
SporeStack Python CLI

Optional environment variables:
SPORESTACK_ENDPOINT
*or*
SPORESTACK_USE_TOR_ENDPOINT

TOR_PROXY (defaults to socks5h://127.0.0.1:9050 which is fine for most)
"""

cli = typer.Typer(help=HELP)

logging.basicConfig(level=logging.INFO)

DEFAULT_FLAVOR = "vps-1vcpu-1gb"

WAITING_PAYMENT_TO_PROCESS = "Waiting for payment to process..."


def get_api_endpoint() -> str:
    api_endpoint = os.getenv("SPORESTACK_ENDPOINT", api_client.CLEARNET_ENDPOINT)
    if os.getenv("SPORESTACK_USE_TOR_ENDPOINT", None) is not None:
        api_endpoint = api_client.TOR_ENDPOINT
    return api_endpoint


def make_payment(currency: str, uri: str, usd: str) -> None:
    import segno

    premessage = """Payment URI: {}
Pay *exactly* the specified amount. No more, no less. Pay within
one hour at the very most.
Resize your terminal and try again if QR code above is not readable.
Press ctrl+c to abort."""
    message = premessage.format(uri)
    qr = segno.make(uri)
    # This typer.echos.
    qr.terminal()
    typer.echo(message)
    typer.echo(f"Approximate price in USD: {usd}")
    input("[Press enter once you have made payment.]")


@cli.command()
def launch(
    hostname: str,
    days: int = typer.Option(...),
    ssh_key_file: Path = typer.Option(...),
    operating_system: str = typer.Option(...),
    flavor: str = DEFAULT_FLAVOR,
    currency: Optional[str] = None,
    settlement_token: Optional[str] = None,
    region: Optional[str] = None,
) -> None:
    """
    Attempts to launch a server.
    """

    from . import utils

    if settlement_token is not None:
        if currency is None or currency == "settlement":
            currency = "settlement"
        else:
            msg = "Cannot use non-settlement --currency with --settlement-token"
            typer.echo(msg, err=True)
            raise typer.Exit(code=2)
    if currency is None:
        typer.echo("--currency must be set.", err=True)
        raise typer.Exit(code=2)

    if machine_exists(hostname):
        typer.echo(f"{hostname} already created.")
        raise typer.Exit(code=1)

    ssh_key = ssh_key_file.read_text()

    machine_id = utils.random_machine_id()

    assert currency is not None
    response = api_client.launch(
        machine_id=machine_id,
        days=days,
        flavor=flavor,
        operating_system=operating_system,
        ssh_key=ssh_key,
        currency=currency,
        region=region,
        settlement_token=settlement_token,
        api_endpoint=get_api_endpoint(),
        retry=True,
    )

    # This will be false at least the first time if paying with BTC or BCH.
    if response.payment.paid is False:
        assert response.payment.uri is not None
        make_payment(
            currency=currency,
            uri=response.payment.uri,
            usd=response.payment.usd,
        )

        tries = 360
        while tries > 0:
            tries = tries - 1
            typer.echo(WAITING_PAYMENT_TO_PROCESS, err=True)
            # FIXME: Wait one hour in a smarter way.
            # Waiting for payment to set in.
            time.sleep(10)
            response = api_client.launch(
                machine_id=machine_id,
                days=days,
                flavor=flavor,
                operating_system=operating_system,
                ssh_key=ssh_key,
                currency=currency,
                region=region,
                settlement_token=settlement_token,
                api_endpoint=get_api_endpoint(),
                retry=True,
            )
            if response.payment.paid is True:
                break

    if response.created is False:
        tries = 360
        while tries > 0:
            typer.echo("Waiting for server to build...", err=True)
            tries = tries + 1
            # Waiting for server to spin up.
            time.sleep(10)
            response = api_client.launch(
                machine_id=machine_id,
                days=days,
                flavor=flavor,
                operating_system=operating_system,
                ssh_key=ssh_key,
                currency=currency,
                region=region,
                settlement_token=settlement_token,
                api_endpoint=get_api_endpoint(),
                retry=True,
            )
            if response.created is True:
                break

    if response.created is False:
        typer.echo("Server creation failed, tries exceeded.", err=True)
        raise typer.Exit(code=1)

    created_dict = response.dict()
    created_dict["vm_hostname"] = hostname
    save_machine_info(created_dict)
    typer.echo(pretty_machine_info(created_dict), err=True)
    typer.echo(json.dumps(created_dict, indent=4))


@cli.command()
def topup(
    hostname: str,
    days: int = typer.Option(...),
    currency: Optional[str] = None,
    settlement_token: Optional[str] = None,
) -> None:
    """
    tops up an existing vm.
    """
    if settlement_token is not None:
        if currency is None or currency == "settlement":
            currency = "settlement"
        else:
            msg = "Cannot use non-settlement --currency with --settlement-token"
            typer.echo(msg, err=True)
            raise typer.Exit(code=2)

    if currency is None:
        typer.echo("--currency must be set.", err=True)
        raise typer.Exit(code=2)

    if not machine_exists(hostname):
        typer.echo(f"{hostname} does not exist.")
        raise typer.Exit(code=1)

    machine_info = get_machine_info(hostname)
    machine_id = machine_info["machine_id"]

    assert currency is not None
    response = api_client.topup(
        machine_id=machine_id,
        days=days,
        currency=currency,
        api_endpoint=get_api_endpoint(),
        settlement_token=settlement_token,
        retry=True,
    )

    # This will be false at least the first time if paying with anything
    # but settlement.
    if response.payment.paid is False:
        assert response.payment.uri is not None
        make_payment(
            currency=currency,
            uri=response.payment.uri,
            usd=response.payment.usd,
        )

        tries = 360
        while tries > 0:
            typer.echo(WAITING_PAYMENT_TO_PROCESS, err=True)
            tries = tries - 1
            # FIXME: Wait one hour in a smarter way.
            # Waiting for payment to set in.
            time.sleep(10)
            response = api_client.topup(
                machine_id=machine_id,
                days=days,
                currency=currency,
                api_endpoint=get_api_endpoint(),
                settlement_token=settlement_token,
                retry=True,
            )
            if response.payment.paid is True:
                break

    machine_info["expiration"] = response.expiration
    save_machine_info(machine_info, overwrite=True)
    typer.echo(machine_info["expiration"])


def server_info_path() -> Path:
    home = os.getenv("HOME")
    assert home is not None, "Unable to detect $HOME environment variable?"
    sporestack_dir = Path(home, ".sporestack")

    # Put servers in a subdirectory
    servers_dir = sporestack_dir.joinpath("servers")

    # Migrate existing server.json files into servers subdirectory
    if sporestack_dir.exists() and not servers_dir.exists():
        typer.echo(
            f"Migrating server profiles found in {sporestack_dir} to {servers_dir}.",
            err=True,
        )
        servers_dir.mkdir()
        for json_file in sporestack_dir.glob("*.json"):
            json_file.rename(servers_dir.joinpath(json_file.name))

    # Make it, if it doesn't exist already.
    sporestack_dir.mkdir(exist_ok=True)
    servers_dir.mkdir(exist_ok=True)

    return servers_dir


def save_machine_info(machine_info: Dict[str, Any], overwrite: bool = False) -> None:
    """
    Save info to disk.
    """
    os.umask(0o0077)
    directory = server_info_path()
    hostname = machine_info["vm_hostname"]
    json_file = directory.joinpath(f"{hostname}.json")
    if overwrite is False:
        assert json_file.exists() is False, f"{json_file} already exists."
    json_file.write_text(json.dumps(machine_info))


def get_machine_info(hostname: str) -> Dict[str, Any]:
    """
    Get info from disk.
    """
    directory = server_info_path()
    json_file = directory.joinpath(f"{hostname}.json")
    if not json_file.exists():
        raise ValueError(f"{hostname} does not exist in {directory} as {json_file}")
    machine_info = json.loads(json_file.read_bytes())
    assert isinstance(machine_info, dict)
    if machine_info["vm_hostname"] != hostname:
        raise ValueError("hostname does not match filename.")
    return machine_info


def pretty_machine_info(info: Dict[str, Any]) -> str:
    msg = "Hostname: {}\n".format(info["vm_hostname"])
    msg += "Machine ID (keep this secret!): {}\n".format(info["machine_id"])
    if "ipv6" in info["network_interfaces"][0]:
        msg += "IPv6: {}\n".format(info["network_interfaces"][0]["ipv6"])
    if "ipv4" in info["network_interfaces"][0]:
        msg += "IPv4: {}\n".format(info["network_interfaces"][0]["ipv4"])
    expiration = info["expiration"]
    human_expiration = time.strftime("%Y-%m-%d %H:%M:%S %z", time.localtime(expiration))
    if "running" in info:
        msg += "Running: {}\n".format(info["running"])
    msg += f"Expiration: {expiration} ({human_expiration})\n"
    time_to_live = expiration - int(time.time())
    hours = time_to_live // 3600
    msg += f"Server will be deleted in {hours} hours."
    return msg


@cli.command()
def list() -> None:
    """
    List all locally known servers.
    """
    from .exceptions import SporeStackUserError

    directory = server_info_path()
    infos = []
    for hostname_json in os.listdir(directory):
        hostname = hostname_json.split(".")[0]
        saved_vm_info = get_machine_info(hostname)
        try:
            upstream_vm_info = api_client.info(
                machine_id=saved_vm_info["machine_id"],
            )
            saved_vm_info["expiration"] = upstream_vm_info.expiration
            saved_vm_info["running"] = upstream_vm_info.running
            infos.append(saved_vm_info)
        except SporeStackUserError as e:
            expiration = saved_vm_info["expiration"]
            human_expiration = time.strftime(
                "%Y-%m-%d %H:%M:%S %z", time.localtime(expiration)
            )
            msg = hostname
            msg += f" expired ({expiration} {human_expiration}): "
            msg += str(e)
            typer.echo(msg)

    for info in infos:
        typer.echo()
        typer.echo(pretty_machine_info(info))

    typer.echo()


def machine_exists(hostname: str) -> bool:
    """
    Check if the VM's JSON exists locally.
    """
    return server_info_path().joinpath(f"{hostname}.json").exists()


@cli.command()
def get_attribute(hostname: str, attribute: str) -> None:
    """
    Returns an attribute about the VM.
    """
    machine_info = get_machine_info(hostname)
    typer.echo(machine_info[attribute])


@cli.command()
def info(hostname: str) -> None:
    """
    Info on the VM
    """
    machine_info = get_machine_info(hostname)
    machine_id = machine_info["machine_id"]
    typer.echo(
        api_client.info(machine_id=machine_id, api_endpoint=get_api_endpoint()).json()
    )


@cli.command()
def start(hostname: str) -> None:
    """
    Boots the VM.
    """
    machine_info = get_machine_info(hostname)
    machine_id = machine_info["machine_id"]
    api_client.start(machine_id=machine_id, api_endpoint=get_api_endpoint())
    typer.echo(f"{hostname} started.")


@cli.command()
def stop(hostname: str) -> None:
    """
    Immediately kills the VM.
    """
    machine_info = get_machine_info(hostname)
    machine_id = machine_info["machine_id"]
    api_client.stop(machine_id=machine_id, api_endpoint=get_api_endpoint())
    typer.echo(f"{hostname} stopped.")


@cli.command()
def delete(hostname: str) -> None:
    """
    Deletes the VM before expiration (no refunds/credits)
    """
    machine_info = get_machine_info(hostname)
    machine_id = machine_info["machine_id"]
    api_client.delete(machine_id=machine_id, api_endpoint=get_api_endpoint())
    # Also remove the .json file
    server_info_path().joinpath(f"{hostname}.json").unlink()
    typer.echo(f"{hostname} was deleted.")


@cli.command()
def rebuild(hostname: str) -> None:
    """
    Rebuilds the VM with the operating system and SSH key given at launch time.

    Will take a couple minutes to complete after the request is made.
    """
    machine_info = get_machine_info(hostname)
    machine_id = machine_info["machine_id"]
    api_client.rebuild(machine_id=machine_id, api_endpoint=get_api_endpoint())
    typer.echo(f"{hostname} rebuilding.")


@cli.command()
def settlement_token_enable(
    token: str,
    dollars: int = typer.Option(...),
    currency: str = typer.Option(...),
) -> None:
    """
    Enables a new settlement token.

    Dollars is starting balance.
    """

    response = api_client.token_enable(
        token=token,
        dollars=dollars,
        currency=currency,
        api_endpoint=get_api_endpoint(),
        retry=True,
    )

    uri = response.payment.uri
    assert uri is not None
    usd = response.payment.usd

    make_payment(currency=currency, uri=uri, usd=usd)

    tries = 360 * 2
    while tries > 0:
        typer.echo(WAITING_PAYMENT_TO_PROCESS, err=True)
        tries = tries - 1
        # FIXME: Wait two hours in a smarter way.
        # Waiting for payment to set in.
        time.sleep(10)
        response = api_client.token_enable(
            token=token,
            dollars=dollars,
            currency=currency,
            api_endpoint=get_api_endpoint(),
            retry=True,
        )
        if response.payment.paid is True:
            typer.echo(
                f"{token} has been enabled with ${dollars}. Save it and don't lose it!"
            )
            return
    raise ValueError(f"{token} did not get enabled in time.")


@cli.command()
def settlement_token_add(
    token: str,
    dollars: int = typer.Option(...),
    currency: str = typer.Option(...),
) -> None:
    """
    Adds balance to an existing settlement token.
    """

    response = api_client.token_add(
        token,
        dollars,
        currency=currency,
        api_endpoint=get_api_endpoint(),
        retry=True,
    )

    uri = response.payment.uri
    assert uri is not None
    usd = response.payment.usd

    make_payment(currency=currency, uri=uri, usd=usd)

    tries = 360 * 2
    while tries > 0:
        typer.echo(WAITING_PAYMENT_TO_PROCESS, err=True)
        tries = tries - 1
        # FIXME: Wait two hours in a smarter way.
        response = api_client.token_add(
            token,
            dollars,
            currency=currency,
            api_endpoint=get_api_endpoint(),
            retry=True,
        )
        # Waiting for payment to set in.
        time.sleep(10)
        if response.payment.paid is True:
            typer.echo(f"Added {dollars} dollars to {token}")
            return
    raise ValueError(f"{token} did not get enabled in time.")


@cli.command()
def settlement_token_balance(token: str) -> None:
    """
    Gets balance for a settlement token.
    """

    typer.echo(
        api_client.token_balance(token=token, api_endpoint=get_api_endpoint()).usd
    )


@cli.command()
def settlement_token_generate() -> None:
    """
    Generates a settlement token that can be enabled.
    """
    from . import utils

    typer.echo(utils.random_token())


@cli.command()
def version() -> None:
    """
    Returns the installed version.
    """
    typer.echo(importlib_metadata_version(__package__))


@cli.command()
def api_endpoint() -> None:
    """
    Prints the selected API endpoint: Env var: SPORESTACK_ENDPOINT,
    or, SPORESTACK_USE_TOR=1
    """
    endpoint = get_api_endpoint()
    if ".onion" in endpoint:
        typer.echo(f"{endpoint} using {api_client._get_tor_proxy()}")
        return
    else:
        typer.echo(endpoint)
        return


if __name__ == "__main__":
    cli()