Coverage for port4me/__init__.py: 86%
101 statements
« prev ^ index » next coverage.py v7.3.3, created at 2023-12-15 17:07 -0800
« prev ^ index » next coverage.py v7.3.3, created at 2023-12-15 17:07 -0800
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
4from itertools import islice
5import socket
6from getpass import getuser
7from os import getenv
10__version__ = "0.6.0-9001"
11__all__ = ["port4me", "port4me_gen"]
14# Source: https://chromium.googlesource.com/chromium/src.git/+/refs/heads/master/net/base/port_util.cc
15# Last updated: 2022-10-24
16unsafe_ports_chrome = getenv("PORT4ME_EXCLUDE_UNSAFE_CHROME", "1,7,9,11,13,15,17,19,20,21,22,23,25,37,42,43,53,69,77,79,87,95,101,102,103,104,109,110,111,113,115,117,119,123,135,137,139,143,161,179,389,427,465,512,513,514,515,526,530,531,532,540,548,554,556,563,587,601,636,989,990,993,995,1719,1720,1723,2049,3659,4045,5060,5061,6000,6566,6665,6666,6667,6668,6669,6697,10080")
18# Source: https://www-archive.mozilla.org/projects/netlib/portbanning#portlist
19# Last updated: 2022-10-24
20unsafe_ports_firefox = getenv("PORT4ME_EXCLUDE_UNSAFE_FIREFOX", "1,7,9,11,13,15,17,19,20,21,22,23,25,37,42,43,53,77,79,87,95,101,102,103,104,109,110,111,113,115,117,119,123,135,139,143,179,389,465,512,513,514,515,526,530,531,532,540,556,563,587,601,636,993,995,2049,4045,6000")
23def uint_hash(s):
24 h = 0
25 for char in s:
26 h = (31 * h + ord(char)) % 2**32
27 return h
30def is_port_free(port):
31 with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
32 try:
33 s.bind(("", port))
34 except PermissionError:
35 return False
36 return True
39def parse_ports(string):
40 ports = []
41 for port in string.replace("{chrome}", unsafe_ports_chrome).replace(
42 "{firefox}", unsafe_ports_firefox).replace(",", " ").split():
43 if port:
44 port1, _, port2 = port.partition("-")
45 if port2:
46 ports.extend(range(int(port1), int(port2)+1))
47 else:
48 ports.append(int(port1))
49 return ports
52def get_env_ports(var_name):
53 """Get an ordered set of ports from the environment variable `var_name` and `var_name`_SITE"""
54 ports = []
55 names = [var_name, var_name+"_SITE"]
56 if var_name == "PORT4ME_EXCLUDE":
57 names.append(var_name+"_UNSAFE")
59 for name in names:
60 if name == "PORT4ME_EXCLUDE_UNSAFE":
61 ports_str = getenv(name, "{chrome},{firefox}")
62 else:
63 ports_str = getenv(name, "")
64 try:
65 ports.extend(parse_ports(ports_str))
66 except ValueError:
67 raise ValueError("invalid port in environment variable "+name)
68 return dict.fromkeys(ports).keys() # discard duplicates but preserve order
71def lcg(seed, a=75, c=74, modulus=65537):
72 """
73 Get the next number in a sequence according to a Linear Congruential Generator algorithm.
75 The default constants are from the ZX81.
76 """
77 seed %= modulus
78 seed_next = (a*seed + c) % modulus
80 # For certain LCG parameter settings, we might end up in the same
81 # LCG state. For example, this can happen when (a-c) = 1 and
82 # seed = modulus-1. To make sure we handle any parameter setup, we
83 # detect this manually, increment the seed, and recalculate.
84 if seed_next == seed:
85 return lcg(seed+1, a, c, modulus)
87 # assert 0 <= seed_next <= modulus
88 return seed_next
91def port4me_gen_unfiltered(tool=None, user=None, prepend=None):
92 if prepend is None:
93 prepend = get_env_ports("PORT4ME_PREPEND")
94 elif isinstance(prepend, str):
95 prepend = parse_ports(prepend)
97 yield from prepend
99 if not user:
100 user = getenv("PORT4ME_USER", getuser())
101 if tool is None:
102 tool = getenv("PORT4ME_TOOL", "")
104 port = uint_hash((user+","+tool).rstrip(","))
105 while True:
106 port = lcg(port)
107 yield port
110def port4me_gen(tool=None, user=None, prepend=None, include=None, exclude=None, min_port=1024, max_port=65535):
111 if include is None:
112 include = get_env_ports("PORT4ME_INCLUDE")
113 elif isinstance(include, str):
114 include = parse_ports(include)
116 if exclude is None:
117 exclude = get_env_ports("PORT4ME_EXCLUDE")
118 elif isinstance(exclude, str):
119 exclude = parse_ports(exclude)
121 for port in port4me_gen_unfiltered(tool, user, prepend):
122 if ((min_port <= port <= max_port)
123 and (not include or port in include)
124 and (not exclude or port not in exclude)):
125 yield port
128_list = list # necessary to avoid conflicts with list() and the parameter which is named list
131def port4me(tool=None, user=None, prepend=None, include=None, exclude=None, skip=None,
132 list=None, test=None, max_tries=65536, must_work=True, min_port=1024, max_port=65535):
133 """
134 Find a free TCP port using a deterministic sequence of ports based on the current username.
136 This reduces the chance of different users trying to access the same port,
137 without having to use a completely random new port every time.
139 Parameters
140 ----------
141 tool : str, optional
142 Used in the seed when generating port numbers, to get a different port sequence for different tools.
143 user : str, optional
144 Used in the seed when generating port numbers. Defaults to determining the username with getuser().
145 prepend : list, optional
146 A list of ports to try first
147 include : list, optional
148 If specified, skip any ports not in this list
149 exclude : list, optional
150 Skip any ports in this list
151 skip : int, optional
152 Skip this many ports at the beginning (after excluded ports have been skipped)
153 list : int, optional
154 Instead of returning a single port, return a list of this many ports without checking if they are free.
155 test : int, optional
156 If specified, return whether the port `test` is not in use. All other parameters will be ignored.
157 max_tries : int, optional
158 Raise a TimeoutError if it takes more than this many tries to find a port. Default is 65536.
159 must_work : bool, optional
160 If True, then an error is produced if no port could be found. If False, then `-1` is returned.
161 min_port : int, optional
162 Skips any ports that are smaller than this
163 max_port : int, optional
164 Skips any ports that are larger than this
165 """
166 if test:
167 return is_port_free(test)
169 tries = 1
171 gen = port4me_gen(tool, user, prepend, include, exclude, min_port, max_port)
173 if skip is None:
174 skip = getenv("PORT4ME_SKIP", 0)
175 skip = int(skip)
176 gen = islice(gen, skip, None)
178 if list is None:
179 list = getenv("PORT4ME_LIST", 0)
180 list = int(list)
182 if list:
183 return _list(islice(gen, list))
185 for port in gen:
186 if is_port_free(port):
187 break
189 if max_tries and tries > max_tries:
190 if must_work:
191 raise TimeoutError("Failed to find a free TCP port after {} attempts".format(max_tries))
192 else:
193 return -1
195 tries += 1
197 return port
200if __name__ == "__main__":
201 print(port4me())