src.wg.pywg

  1#!/usr/bin/python3
  2
  3import argparse
  4import functools
  5import logging
  6import pathlib
  7import platform
  8from ipaddress import IPv4Interface
  9from subprocess import Popen
 10from time import time
 11
 12if platform.system() == "Linux":
 13    import fcntl
 14
 15logger = logging.getLogger()
 16
 17
 18WIREGUARD_CONF = "/etc/wireguard/wg1.conf"
 19# WIREGUARD_CONF = "C:\\code\\vpn_dan_bot\\src\\wg\\wg1.conf"
 20WIREGUARD_LOG = "/home/zadira/Scripts/ErrorsPyWG"
 21
 22
 23class ConfigChanger:
 24    def __init__(self, public_key, allowed_ips=None, raises=False) -> None:
 25        self.public_key = public_key
 26        self.allowed_ips = allowed_ips
 27        self.raises = raises
 28
 29    def __find_key(self):
 30        lines = []
 31        for i, line in enumerate(self.config):
 32            if self.public_key in line:
 33                lines.append(i)
 34
 35        if not len(lines):
 36            raise Exception("PEER NOT FOUND")
 37
 38        validated_lines = []
 39        for i in lines:
 40            assert "peer" in self.config[i - 1].lower()
 41            validated_lines.extend((i - 1, i))
 42
 43            extra = 1
 44            while extra:
 45                try:
 46                    if (
 47                        "peer" in self.config[i + extra].lower()
 48                        or not self.config[i + extra]
 49                    ):
 50                        break
 51                except IndexError:
 52                    break
 53                else:
 54                    validated_lines.append(i + extra)
 55                    extra += 1
 56
 57        for i in validated_lines:
 58            yield i
 59
 60    def process(mode):
 61        def decorator(func):
 62            @functools.wraps(func)
 63            def conf_changer(self, *args, **kwargs):
 64                try:
 65                    start = time()
 66
 67                    with open(WIREGUARD_CONF, "r+") as file:
 68                        fcntl.flock(file, fcntl.LOCK_EX)
 69                        self.config = file.read().split("\n")
 70                        func(self, *args, **kwargs)
 71
 72                        file.seek(0)
 73                        file.truncate(0)
 74                        file.write("\n".join(self.config))
 75                        fcntl.flock(file, fcntl.LOCK_UN)
 76
 77                    proc = time() - start
 78
 79                    reload_cmd = " && ".join(
 80                        (
 81                            f'echo -n "$(date) :: {func.__name__} :: {self.public_key[:4]}... ::" >> {WIREGUARD_LOG}',
 82                            f"sudo systemctl reload wg-quick@wg1.service 2>> {WIREGUARD_LOG}",
 83                            f'echo " " >> {WIREGUARD_LOG}',
 84                        )
 85                    )
 86
 87                    error_code = Popen(
 88                        f"flock {WIREGUARD_LOG} --command '{reload_cmd}'", shell=True
 89                    ).returncode
 90                    if not error_code:
 91                        logger.info(f'Peer "{self.public_key[:4]}..." {mode}ed')
 92                    else:
 93                        raise Exception(
 94                            f"execute wireguard reload error. Exit code {error_code}"
 95                        )
 96
 97                except Exception as e:
 98                    logger.error(e.args[0])
 99                    if self.raises:
100                        raise e
101                else:
102                    logger.debug(
103                        f"METRIC::{func.__name__}::proc::[  {int(proc)*1000}  ]msec"
104                    )
105                finally:
106                    logger.debug(
107                        f"METRIC::{func.__name__}::end::[  {int((time()-start)*1000)}  ]msec"
108                    )
109
110            return conf_changer
111
112        return decorator
113
114    @process("add")
115    def register(self):
116        new_user = (
117            "[Peer]",
118            f"PublicKey = {self.public_key}",
119            f"AllowedIPs = {self.allowed_ips}",
120            "PersistentKeepalive = 25",
121        )
122        config = "\n".join(self.config)
123        if self.public_key in config:
124            logger.warning("Peer already added")
125            raise Exception("NEW USER ERROR")
126
127        self.config.extend(new_user)
128
129    @process("delet")
130    def delete(self):
131        lines = self.__find_key()
132        for n in sorted(lines, reverse=True):
133            self.config.pop(n)
134
135    @process("bann")
136    def ban(self):
137        lines = self.__find_key()
138        for n in lines:
139            if self.config[n].startswith("#"):
140                logger.warning("Peer already banned")
141                raise Exception("BAN USER ERROR")
142            else:
143                self.config[n] = f"# {self.config[n]}"
144
145    @process("unbann")
146    def unban(self):
147        lines = self.__find_key()
148        for n in lines:
149            if not self.config[n].startswith("#"):
150                logger.warning("Peer already unbanned")
151                raise Exception("UNBAN USER ERROR")
152            else:
153                self.config[n] = self.config[n].strip("# ")
154
155
156def parse():
157    parser = argparse.ArgumentParser(
158        prog="Python Wireguard Utils",
159        description="Python tools for the wireguard server",
160        epilog="Text at the bottom of help",
161    )
162
163    parser.add_argument(
164        "pubkey", type=str, help="Public key of new(old) user", nargs="?"
165    )
166    parser.add_argument(
167        "-ips", "--allowed_ips", type=IPv4Interface, help="AllowedIPs for new peer"
168    )
169    parser.add_argument(
170        "-m",
171        "--mode",
172        choices=["new", "ban", "unban", "del"],
173        default="new",
174        help="Tool for change wireguard config (default 'new')",
175    )
176    parser.add_argument(
177        "--wgpath",
178        type=pathlib.Path,
179        default=WIREGUARD_CONF,
180        help=f"Path wireguard config file (default '{WIREGUARD_CONF}')",
181    )
182    parser.add_argument(
183        "--raises",
184        action="store_true",
185        help="Interruption in case of an error",
186    )
187    parser.add_argument(
188        "-l",
189        "--list",
190        action="store_true",
191        help="Print list of peers (banned too)",
192    )
193    args = parser.parse_args()
194
195    return args
196
197
198def main():
199    args = parse()
200
201    if args.list:
202        with open(WIREGUARD_CONF) as file:
203            peers = file.read().split("[Peer]")[1:]
204            for peer in peers:
205                print("\x1b[33m[Peer]\x1b[0m", end=" ")
206
207                if "\n#" in peer.strip("#\n "):
208                    print("\x1b[31;1m(BAN)\x1b[0m")
209                else:
210                    print("")
211
212                print(peer.strip(" #\n").replace("# ", ""))
213
214        return
215
216    cc = ConfigChanger(
217        public_key=args.pubkey, allowed_ips=args.allowed_ips, raises=args.raises
218    )
219    if args.mode == "new":
220        if not args.pubkey:
221            logger.error("Argument PUBKEY: empty value")
222        elif not args.allowed_ips:
223            logger.error("Argument ALLOWED_IPS: empty value")
224        else:
225            cc.register()
226    elif args.mode == "ban":
227        cc.ban()
228    elif args.mode == "unban":
229        cc.unban()
230    elif args.mode == "del":
231        cc.delete()
232
233
234if __name__ == "__main__":
235    logging.basicConfig(
236        level=logging.DEBUG,
237        format="[%(levelname)-8s] %(funcName)s(%(lineno)d): %(message)s",
238    )
239    logger.info("Start")
240    main()
241    logger.info("End")