diff --git a/proton/vpn/cli/__init__.py b/proton/vpn/cli/__init__.py index 002c7bf..eb5383d 100644 --- a/proton/vpn/cli/__init__.py +++ b/proton/vpn/cli/__init__.py @@ -29,7 +29,7 @@ from dbus_fast import BusType, Message, MessageType from proton.vpn.cli.commands.account import signin, signout, info -from proton.vpn.cli.commands.server import connect, disconnect +from proton.vpn.cli.commands.server import connect, disconnect, status from proton.vpn.cli.commands.location_discovery import countries, cities from proton.vpn.cli.commands.settings import config from proton.vpn.cli.core.controller import Params @@ -121,6 +121,7 @@ async def app(ctx, verbose): # server related functionality app.add_command(connect) +app.add_command(status) app.add_command(disconnect) # listing functionality diff --git a/proton/vpn/cli/commands/server.py b/proton/vpn/cli/commands/server.py index ec9aad2..11eca4e 100644 --- a/proton/vpn/cli/commands/server.py +++ b/proton/vpn/cli/commands/server.py @@ -19,10 +19,12 @@ You should have received a copy of the GNU General Public License along with ProtonVPN. If not, see . """ -from asyncio import CancelledError -from typing import Optional import click +from json import loads, dumps +from pathlib import Path +from asyncio import CancelledError +from typing import Optional from proton.vpn.cli.core.exceptions import \ AuthenticationRequiredError, \ @@ -35,10 +37,14 @@ from proton.vpn.session.exceptions import ServerNotFoundError from proton.vpn.session.servers.types import LogicalServer, ServerFeatureEnum from proton.vpn.cli.commands.account import SIGNIN_COMMAND +from proton.vpn.connection import events from proton.vpn.cli.commands.command_utils import \ inform_that_expired_serverlist_will_be_updated_if_necessary +TMP_FILE = Path("/tmp/protonvpn-connection_details.json") + + class FailedConnection(click.ClickException): """When attempting to establish a connection, it fails """ @@ -150,6 +156,11 @@ async def connect( ) if server_ip: click.echo(f"Your new IP address is {server_ip}.") + TMP_FILE.write_text(dumps({ + key: getattr(connection_details, key, "") + for key in dir(connection_details) + if not key.startswith("_") + }), encoding="utf-8") protocol = (await controller.get_settings()).protocol _display_openvpn_warning_if_necessary(protocol) @@ -160,9 +171,93 @@ async def connect( "Try connecting to a different server or check your network settings." ) + CONNECT_COMMAND = connect.name +@click.command() +@click.pass_context +@click.option("-j", "--json", is_flag=True, help="Dump status as JSON") +@click.option("-s", "--simple", is_flag=True, help="Print simple status to stdout, useful for i3bar or polybar.") +@run_async +async def status(ctx, json: bool = False, simple: bool = False): + """Get status from Proton VPN""" + controller = await Controller.create(params=ctx.obj, click_ctx=ctx) + connector = await controller.get_vpn_connector() + connection = connector.current_connection + + # Bail early if disconnected. + if connection is None: + click.echo("ProtonVPN: Disconnected") + return + + server = await controller.find_logical_server(connection.server_name) + state = connector.current_state + + # Bail early if disconnected. + if state is None: + click.echo("ProtonVPN: Disconnected") + return + + # _Try_ to get connection details... sometimes it's not populated though??? + connection_details = state.context.event.context.connection_details + + if connection_details and not TMP_FILE.exists(): + TMP_FILE.write_text(dumps({ + key: getattr(connection_details, key, "") + for key in dir(connection_details) + if not key.startswith("_") + }), encoding="utf-8") + + # If it's not populated, read from tmp file on disk, if we can + try: + connection_details = connection_details or events.ConnectionDetails( + **loads(TMP_FILE.read_text(encoding="utf-8")) + ) + except OSError, TypeError: + # Either the file wasn't there, permissions were bad, or the json was not in an expected format / malformed + pass + + status = { + "exit": { + "city": server.city, + "country": server.exit_country, + "name": connection.server_name, + "ipv4": getattr(connection_details, "server_ipv4", ""), + "ipv6": getattr(connection_details, "server_ipv6", "") + }, + "entry": { + "country": server.entry_country, + "ip": connection.server_ip, + "name": connection.server_name, + "client": { + "ip": getattr(connection_details, "device_ip", ""), + "country": getattr(connection_details, "device_country", "") + } + }, + } + + if json: + return dumps(status) + + elif simple: + _ip = status["exit"].get("ipv4", "") + click.echo( + f"ProtonVPN: {status["exit"]["name"]}" + f"{f" ({_ip})" if _ip else ""}" + ) + return status + + else: + click.echo( + f"Connected to {status["exit"]["name"]} " + f"in {_get_most_specific_server_location(server)}. " + f"Your new IP address is {status["exit"].get("ipv4", "Unknown")}." + ) + # print(status) + return status + + @click.command() @click.pass_context @run_async