wg-vanity.py

#!/usr/bin/env python
"""Prefix wrapper for wireguard-vanity-address

* Improve performance for generating keys with common prefixes.
* Extract prefixes delimited by a '+' or '/'.
* Exit at least on match for each prefix.

Usage:
    $ wg-vanity.py ex{1,2,3}
    Searching for ['ex1/', 'ex1+', 'ex2/', 'ex2+', 'ex3/', 'ex3+'] with ['wireguard-vanity-address', '--in', '2', 'ex']
    private 0AqLDsMHNBFtB1BDr5QFm4un0wmKfCXUm4rWtWjwW2E=  public EX3+ArEAdcwj7xU92idso00qIAK66RQtQNJYJEJOCVY=
    private aBc3/BY3yBahqBLEjv8Bu5LGyR6XywSOFOrFdKaHb3w=  public ex2+BsIonJ3gtDuuFy2BGYdEm0lSlt2Gh/EtgwAs7U8=
    private AFT9ZAZachKiH0Oig9s6qFKrDuWevItYmDnU4QkMp08=  public eX3+aoyOahwll9ppovoMbC9dFnSImsgPWhUDJ5yfVBg=
    private AMO0Il76vI1PbaBEjzcL9wwD0yM/BxS0T7CvQy6kU08=  public EX1//zq/EZkD4qs42zGaP0yH4bW4mRcI2MORGVFZZTc=

TODO:
* Replace this script with proper rust fix that searches for multiplex prefixes
  at the same time.
"""

import os
import subprocess
import sys

def find_common_prefix_and_unique_components(str_list: list[str]):
    # Find the common prefix
    common_prefix = ""
    for i, char_group in enumerate(zip(*str_list)):
        if len(set(char_group)) == 1:
            common_prefix += char_group[0]
        else:
            break

    # Extract the unique components
    unique_components = [s[len(common_prefix):] for s in str_list]

    return common_prefix, unique_components

def run_wireguard_vanity_address(name: str, matches: int = 1):
    common, uniq = find_common_prefix_and_unique_components(name)
    common = common.lower()
    cmd = ['wireguard-vanity-address', '--in', str(len(common)), common]

    prefixes = []
    prefix_cnt = {}

    for u in uniq:
        prefix = f"{common}{u}"
        prefix_cnt[prefix] = 0
        prefixes += [f"{prefix}/", f"{prefix}+"]

    print(f"Searching for {prefixes} with {cmd}")

    # Start and end of the public key substring we care about
    start = 61
    end = start + len(common) + max(len(u) for u in uniq) + 1

    with subprocess.Popen(cmd, stdout=subprocess.PIPE, text=True, bufsize=1) as p:
        for line in iter(p.stdout.readline, ''):
            prefix = line[start:end].lower()

            if any(m == prefix for m in prefixes):
                print(line.strip())

                prefix_cnt[prefix[:-1]] += 1

                if all(cnt >= matches for cnt in prefix_cnt.values()):
                    break

        os.kill(p.pid, 15)

def main():
    import argparse

    # Create ArgumentParser object
    parser = argparse.ArgumentParser(description='Prefix and match wrapper for wireguard-vanity-address')
    parser.add_argument('--matches', '-m', type=int, default=1, help='Minumium matches per prefix')
    parser.add_argument('prefixes', nargs='+', help='Prefixes to search for')
    args = parser.parse_args()

    run_wireguard_vanity_address(args.prefixes, args.matches)

if __name__ == "__main__":
    main()