Coverage for port4me/__init__.py: 86%

1#!/usr/bin/env python 

2# -*- coding: utf-8 -*- 

3 

4from itertools import islice 

5import socket 

6from getpass import getuser 

7from os import getenv 

8 

9 

10__version__ = "0.6.0-9001" 

11__all__ = ["port4me", "port4me_gen"] 

12 

13 

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") 

17 

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") 

21 

22 

23def uint_hash(s): 

24 h = 0 

25 for char in s: 

26 h = (31 * h + ord(char)) % 2**32 

27 return h 

28 

29 

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 

37 

38 

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 

50 

51 

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") 

58 

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 

69 

70 

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. 

74 

75 The default constants are from the ZX81. 

76 """ 

77 seed %= modulus 

78 seed_next = (a*seed + c) % modulus 

79 

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) 

86 

87 # assert 0 <= seed_next <= modulus 

88 return seed_next 

89 

90 

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) 

96 

97 yield from prepend 

98 

99 if not user: 

100 user = getenv("PORT4ME_USER", getuser()) 

101 if tool is None: 

102 tool = getenv("PORT4ME_TOOL", "") 

103 

104 port = uint_hash((user+","+tool).rstrip(",")) 

105 while True: 

106 port = lcg(port) 

107 yield port 

108 

109 

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) 

115 

116 if exclude is None: 

117 exclude = get_env_ports("PORT4ME_EXCLUDE") 

118 elif isinstance(exclude, str): 

119 exclude = parse_ports(exclude) 

120 

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 

126 

127 

128_list = list # necessary to avoid conflicts with list() and the parameter which is named list 

129 

130 

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. 

135 

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. 

138 

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) 

168 

169 tries = 1 

170 

171 gen = port4me_gen(tool, user, prepend, include, exclude, min_port, max_port) 

172 

173 if skip is None: 

174 skip = getenv("PORT4ME_SKIP", 0) 

175 skip = int(skip) 

176 gen = islice(gen, skip, None) 

177 

178 if list is None: 

179 list = getenv("PORT4ME_LIST", 0) 

180 list = int(list) 

181 

182 if list: 

183 return _list(islice(gen, list)) 

184 

185 for port in gen: 

186 if is_port_free(port): 

187 break 

188 

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 

194 

195 tries += 1 

196 

197 return port 

198 

199 

200if __name__ == "__main__": 

201 print(port4me())