protonne.proton

  1from dataclasses import dataclass
  2from pathlib import Path
  3
  4from morbin import Morbin, Output
  5
  6
  7@dataclass
  8class Server:
  9    name: str
 10    country: str
 11    protocol: str
 12    load: str
 13    plan: str
 14    features: str | None
 15
 16
 17@dataclass(init=False)
 18class KillSwitch:
 19    active: bool
 20    on: bool
 21    permanent: bool
 22
 23    def __init__(self, status: str):
 24        if status == "Off":
 25            self.on = False
 26            self.active = False
 27            self.permanent = False
 28        elif status.startswith("On (Inactive,"):
 29            self.on = True
 30            self.active = False
 31            self.permanent = False
 32        elif status == "On":
 33            self.on = True
 34            self.active = True
 35            self.permanent = False
 36        elif status == "Permanent":
 37            self.on = True
 38            self.active = True
 39            self.permanent = True
 40
 41
 42@dataclass(init=False)
 43class Connection:
 44    IP: str
 45    killswitch: KillSwitch
 46    raw: str
 47    server: Server
 48    time: str
 49
 50    def __init__(self, status: str):
 51        self.raw = status
 52        lines = [line for line in status.splitlines() if ":" in line]
 53        connection = {
 54            line.split(":", 1)[0]: line.split(":", 1)[1].replace("\t", "").strip()
 55            for line in lines
 56        }
 57        self.IP = connection["IP"]
 58        self.server = Server(
 59            connection["Server"],
 60            connection["Country"],
 61            connection["Protocol"],
 62            connection["Server Load"],
 63            connection["Server Plan"],
 64            connection.get("Server Features"),
 65        )
 66        self.killswitch = KillSwitch(connection["Kill switch"])
 67        self.time = connection["Connection time"]
 68
 69    def __str__(self) -> str:
 70        return self.raw
 71
 72
 73class Proton(Morbin):
 74    @property
 75    def program(self) -> str:
 76        return "protonvpn-cli"
 77
 78    # Seat |=============================== Core ===============================|
 79
 80    def config(self, args: str = "") -> Output:
 81        return self.run(f"config {args}")
 82
 83    def connect(
 84        self,
 85        server: str = "",
 86        protocol: str = "",
 87        country_code: str = "",
 88        fastest: bool = False,
 89        random: bool = False,
 90        secure_core: bool = False,
 91        p2p: bool = False,
 92        tor: bool = False,
 93    ) -> Output:
 94        """
 95        #### :params:
 96
 97        `server`: Connect to this specific server (e.g. `CH#4`, `CH-US-1`, `HK5-Tor`).
 98
 99        `protocol`: Connect using this protocol (`tcp` or `udp`).
100
101        `country_code`: Connect to a server in this country (`SE`, `PT`, `BR`, `AR`, etc.).
102
103        `fastest`: Connect to the fastest server.
104
105        `random`: Connect to a random server.
106
107        `secure_core`: Connect to the fastest Secure-Core server.
108
109        `p2p`: Connect to the fastest P2P server.
110
111        `tor`: Connect to the fastest Tor server.
112
113        Supplying an argument for `server` will override other parameters except `protocol`.
114
115        If supplying one of the `bool` params (`fastest`, `random`, `secure_core`, `p2p`, `tor`), supply only one.
116        They will be checked in the order they appear in the function signature and the first that is `True` will be used.
117        """
118        args = ""
119        if server:
120            args = server
121        elif fastest:
122            args = "--fastest"
123        elif random:
124            args = "--random"
125        elif secure_core:
126            args = "--sc"
127        elif p2p:
128            args = "--p2p"
129        elif tor:
130            args = "--tor"
131        if country_code and (not server):
132            args += f" --cc {country_code}"
133        if protocol:
134            args += f" --protocol {protocol}"
135        return self.run(f"connect {args}")
136
137    def disconnect(self) -> Output:
138        return self.run("disconnect")
139
140    def killswitch(self, arg: str) -> Output:
141        """`arg` should be one of `on`, `off`, or `permanent`."""
142        return self.run(f"killswitch --{arg}")
143
144    def login(self, username: str) -> Output:
145        return self.run(f"login {username}")
146
147    def logout(self) -> Output:
148        return self.run("logout")
149
150    def netshield(self, args: str = "") -> Output:
151        return self.run(f"netshield {args}")
152
153    def reconnect(self) -> Output:
154        return self.run("reconnect")
155
156    def status(self) -> Output:
157        """Execute status command."""
158        return self.run("status")
159
160    # Seat |=========================== Convenience ===========================|
161
162    @property
163    def connected(self) -> bool:
164        """Returns whether this device is connected to Proton."""
165        with self.capturing_output():
166            status = self.status().stdout
167        return "Proton VPN Connection Status" in status
168
169    @property
170    def connection(self) -> Connection | None:
171        """If this device is connected, a `Connection` object will be returned.
172        If disconnected, `None` will be returned.
173
174        Accessing this property can be time consuming due to the protonvpn-cli backend.
175        Ideally store it in a local variable until you need to check for an updated connection status.
176        """
177        with self.capturing_output():
178            if self.connected:
179                return Connection(self.status().stdout)
180            return None
181
182    def clear_cache(self):
183        """Clears files from `home/.cache/protonvpn`.
184
185        This seems to help with the time taken to connect sometimes,
186        but if the permanent kill switch is engaged and you try to connect after clearing the cache, you will get an error.
187        """
188        path = Path.home() / ".cache" / "protonvpn"
189        if path.exists():
190            [path.unlink() for path in path.glob("*.*")]
191
192    def connect_fastest(self) -> Output:
193        """Connect to the fastest server."""
194        return self.connect(fastest=True)
195
196    def connect_random(self) -> Output:
197        """Connect to a random server."""
198        return self.connect(random=True)
199
200    def disable_killswitch(self) -> Output:
201        return self.killswitch("off")
202
203    def enable_killswitch(self) -> Output:
204        return self.killswitch("on")
205
206    def enable_permanent_killswitch(self) -> Output:
207        return self.killswitch("permanent")