Package wtf :: Module osutil
[hide private]
[frames] | no frames]

Source Code for Module wtf.osutil

  1  # -*- coding: ascii -*- 
  2  # 
  3  # Copyright 2006-2012 
  4  # Andr\xe9 Malo or his licensors, as applicable 
  5  # 
  6  # Licensed under the Apache License, Version 2.0 (the "License"); 
  7  # you may not use this file except in compliance with the License. 
  8  # You may obtain a copy of the License at 
  9  # 
 10  #     http://www.apache.org/licenses/LICENSE-2.0 
 11  # 
 12  # Unless required by applicable law or agreed to in writing, software 
 13  # distributed under the License is distributed on an "AS IS" BASIS, 
 14  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
 15  # See the License for the specific language governing permissions and 
 16  # limitations under the License. 
 17  """ 
 18  OS Specific Utilities 
 19  ===================== 
 20   
 21  Certain utilities to make the life more easy. 
 22  """ 
 23  __author__ = u"Andr\xe9 Malo" 
 24  __docformat__ = "restructuredtext en" 
 25   
 26  import datetime as _datetime 
 27  import errno as _errno 
 28  import fcntl as _fcntl 
 29  import os as _os 
 30  import resource as _resource 
 31  import socket as _socket 
 32  import sys as _sys 
 33  import threading as _threading 
 34  import warnings as _warnings 
 35   
 36  from wtf import Error, WtfWarning 
 37   
 38   
39 -class IdentityWarning(WtfWarning):
40 """ The attempt to change identity caused a soft error """
41
42 -class IdentityError(Error):
43 """ The attempt to change identity caused a hard error """
44 45
46 -class SocketError(Error):
47 """ Socket error """
48
49 -class AddressError(SocketError):
50 """ Address resolution error """
51
52 -class TimeoutError(SocketError):
53 """ Timeout error """
54
55 -class SSLError(SocketError):
56 """ SSL error """
57 58
59 -def raise_socket_error(timeout=None):
60 """ 61 Convert a socket error into an appropriate module exception 62 63 This function needs an already raised ``socket.error``. 64 65 ``raise_socket_error.EAIS`` is a mapping from GAI error numbers to their 66 names (``{int: 'name', ...}``) 67 68 :Parameters: 69 - `timeout`: applied timeout in seconds, used for the TimeoutError 70 description 71 72 :Types: 73 - `timeout`: ``float`` 74 75 :Exceptions: 76 - `TimeoutError`: ``socket.timeout`` 77 - `AddressError`: address/host resolution error 78 (``socket.gaierror/herror``) 79 - `SSLError`: ``socket.sslerror`` 80 - `SocketError`: other socket errors, ``IOError`` 81 - `Exception`: unrecognized exceptions 82 """ 83 try: 84 raise 85 86 except _socket.timeout: 87 if timeout is not None: 88 raise TimeoutError, "Timed out after %s seconds" % timeout, \ 89 _sys.exc_info()[2] 90 raise TimeoutError, "Timed out", _sys.exc_info()[2] 91 92 except _socket.gaierror, e: 93 # pylint: disable = E1101 94 raise AddressError, "Address Information Error: %s (%s)" % \ 95 (raise_socket_error.EAIS.get(e[0], e[0]), e[1]), \ 96 _sys.exc_info()[2] 97 98 except _socket.herror, e: 99 raise AddressError, "Host Resolution Error %s: %s" % \ 100 (e[0], e[1]), _sys.exc_info()[2] 101 102 except _socket.sslerror, e: 103 raise SSLError, "Socket SSL Error: %s" % str(e), _sys.exc_info()[2] 104 105 except _socket.error, e: 106 if len(e.args) == 1: 107 raise SocketError, "Socket Error: %s" % \ 108 (e[0],), _sys.exc_info()[2] 109 else: 110 raise SocketError, "Socket Error %s: %s" % \ 111 (_errno.errorcode.get(e[0], e[0]), e[1]), _sys.exc_info()[2] 112 113 except IOError, e: 114 raise SocketError, "Socket Error %s: %s" % \ 115 (_errno.errorcode.get(e[0], e[0]), str(e)), \ 116 _sys.exc_info()[2]
117 118 if 1: 119 raise_socket_error.EAIS = dict((val, var) # pylint: disable = W0612 120 for var, val in vars(_socket).items() if var.startswith('EAI_') 121 ) 122 123 139 140
141 -def close_on_exec(descriptor, close=True):
142 """ 143 Mark `descriptor` to be closed on exec (or not) 144 145 :Warning: This function is not thread safe (race condition) 146 147 :Parameters: 148 - `descriptor`: An object with ``fileno`` method or an ``int`` 149 representing a low level file descriptor 150 - `close`: Mark being closed on exec? 151 152 :Types: 153 - `descriptor`: ``file`` or ``int`` 154 - `close`: ``bool`` 155 156 :Exceptions: 157 - `IOError`: Something went wrong 158 """ 159 try: 160 fileno = descriptor.fileno 161 except AttributeError: 162 fd = descriptor 163 else: 164 fd = fileno() 165 166 old = _fcntl.fcntl(fd, _fcntl.F_GETFD) 167 if close: 168 new = old | _fcntl.FD_CLOEXEC 169 else: 170 new = old & ~_fcntl.FD_CLOEXEC 171 _fcntl.fcntl(fd, _fcntl.F_SETFD, new)
172 173
174 -def safe_fd(fd):
175 """ 176 Ensure that file descriptor fd is >= 3 177 178 This is done by dup(2) calls until it's greater than 2. After success 179 the duped descriptors are closed. 180 181 :Parameters: 182 - `fd`: The file descriptor to process 183 184 :Types: 185 - `fd`: ``int`` 186 187 :return: The new file descriptor (>=3) 188 :rtype: ``int`` 189 190 :Exceptions: 191 - `OSError`: Duping went wrong 192 """ 193 toclose = [] 194 try: 195 while fd < 3: 196 toclose.append(fd) 197 fd = _os.dup(fd) 198 finally: 199 for dfd in toclose: 200 try: 201 _os.close(dfd) 202 except OSError: 203 pass 204 return fd
205 206
207 -def close_descriptors(*keep):
208 """ Close all file descriptors >= 3 """ 209 keep = set(keep) 210 try: 211 flag = _resource.RLIMIT_NOFILE 212 except AttributeError: 213 try: 214 flag = _resource.RLIMIT_OFILE 215 except AttributeError: 216 flag = None 217 if flag is not None: 218 try: 219 maxfiles = _resource.getrlimit(flag)[0] 220 except (_resource.error, ValueError): 221 flag = None 222 if flag is None: 223 maxfiles = 256 # wild guess 224 for fd in xrange(3, maxfiles + 1): 225 if fd in keep: 226 continue 227 try: 228 _os.close(fd) 229 except OSError: 230 pass
231 232 233 try: 234 _myflag = _socket.TCP_NODELAY 235 except AttributeError:
236 - def disable_nagle(sock, peername=None):
237 """ 238 Disable nagle algorithm for a TCP socket 239 240 :Note: This function is a NOOP on this platform (not implemented). 241 242 :Parameters: 243 - `sock`: Socket to process 244 - `peername`: The name of the remote socket, if ``str``, it's a UNIX 245 domain socket and the function does nothing 246 247 :Types: 248 - `sock`: ``socket.socket`` 249 - `peername`: ``str`` or ``tuple`` 250 251 :return: The socket and the peername again (if the latter was passed 252 as ``None``, it will be set to something useful 253 :rtype: ``tuple`` 254 255 :Exceptions: 256 - `socket.error`: The socket was probably not connected. If setting 257 of the option fails, no socket error is thrown though. It's ignored. 258 """ 259 if peername is None: 260 peername = sock.getpeername() 261 return sock, peername
262 else:
263 - def disable_nagle(sock, peername=None, _flag=_myflag):
264 """ 265 Disable nagle algorithm for a TCP socket 266 267 :Parameters: 268 - `sock`: Socket to process 269 - `peername`: The name of the remote socket, if ``str``, it's a UNIX 270 domain socket and the function does nothing 271 272 :Types: 273 - `sock`: ``socket.socket`` 274 - `peername`: ``str`` or ``tuple`` 275 276 :return: The socket and the peername again (if the latter was passed 277 as ``None``, it will be set to something useful 278 :rtype: ``tuple`` 279 280 :Exceptions: 281 - `socket.error`: The socket was probably not connected. If setting 282 of the option fails, no socket error is thrown though. It's ignored. 283 """ 284 if peername is None: 285 peername = sock.getpeername() 286 if not isinstance(peername, str): 287 try: 288 sock.setsockopt(_socket.IPPROTO_TCP, _flag, 1) 289 except _socket.error: 290 pass # would have been nice, but, well, not that critical 291 return sock, peername
292 293 294 _connect_cache = {} 295 _connect_cache_lock = _threading.Lock()
296 -def connect(spec, timeout=None, nagle_off=True, cache=0, 297 _cache=_connect_cache, _lock=_connect_cache_lock):
298 """ 299 Create and connect a socket to a peer 300 301 :Parameters: 302 - `spec`: The peer specification (``(host, port)`` or ``str``) 303 - `timeout`: Timeout in seconds 304 - `nagle_off`: Disable Nagle's algorithm. This option does not 305 apply to UNIX domain sockets. 306 307 :Types: 308 - `spec`: ``tuple`` or ``str`` 309 - `timeout`: ``float`` 310 - `nagle_off`: ``bool`` 311 312 :return: The connected socket or ``None`` if no connectable address 313 could be found 314 :rtype: ``socket.socket`` 315 316 :Exceptions: 317 - `SocketError`: socket error (maybe a subclass of `SocketError`) 318 - `NotImplementedError`: UNIX domain sockets are not supported in this 319 platform 320 """ 321 # pylint: disable = W0102, R0912, R0915 322 323 sock = None 324 try: 325 adi = None 326 if cache > 0: 327 _lock.acquire() 328 try: 329 if spec in _cache: 330 adi, stamp = _cache[spec] 331 if stamp < _datetime.datetime.utcnow(): 332 del _cache[spec] 333 adi = None 334 finally: 335 _lock.release() 336 if adi is None: 337 if isinstance(spec, str): 338 try: 339 AF_UNIX = _socket.AF_UNIX 340 except AttributeError: 341 raise NotImplementedError( 342 "UNIX domain sockets are not supported" 343 ) 344 adi = [(AF_UNIX, _socket.SOCK_STREAM, 0, None, spec)] 345 else: 346 adi = _socket.getaddrinfo(spec[0], spec[1], 347 _socket.AF_UNSPEC, _socket.SOCK_STREAM, 0, 0) 348 if cache > 0: 349 _lock.acquire() 350 try: 351 if spec not in _cache: 352 _cache[spec] = ( 353 adi, 354 _datetime.datetime.utcnow() 355 + _datetime.timedelta(seconds=cache), 356 ) 357 finally: 358 _lock.release() 359 360 AF_INET6 = getattr(_socket, 'AF_INET6', None) 361 for family, stype, proto, _, addr in adi: 362 if not _socket.has_ipv6 and family == AF_INET6: 363 continue # skip silenty if python was built without it. 364 365 sock = _socket.socket(family, stype, proto) 366 sock.settimeout(timeout) 367 retry = True 368 while retry: 369 try: 370 sock.connect(addr) 371 except _socket.timeout: 372 break 373 except _socket.error, e: 374 if e[0] == _errno.EINTR: 375 continue 376 elif e[0] in (_errno.ENETUNREACH, _errno.ECONNREFUSED): 377 break 378 raise 379 retry = False 380 else: 381 if nagle_off: 382 disable_nagle(sock) 383 return sock 384 sock.close() 385 except (_socket.error, IOError): 386 try: 387 raise_socket_error(timeout=timeout) 388 except SocketError: 389 e = _sys.exc_info() 390 try: 391 if sock is not None: 392 sock.close() 393 finally: 394 try: 395 raise e[0], e[1], e[2] 396 finally: 397 del e 398 return None
399 400 del _connect_cache, _connect_cache_lock 401 402
403 -def change_identity(user, group):
404 """ 405 Change identity of the current process 406 407 This only works if the effective user ID of the current process is 0. 408 409 :Parameters: 410 - `user`: User identification, if it is interpretable as ``int``, it's 411 assumed to be a numeric user ID 412 - `group`: Group identification, if it is interpretable as ``int``, it's 413 asummed to be a numeric group ID 414 415 :Types: 416 - `user`: ``str`` 417 - `group`: ``str`` 418 419 :Exceptions: 420 - `IdentityWarning`: A soft error occured (like not being root) 421 """ 422 if _os.geteuid() != 0: 423 _warnings.warn("Not attempting to change identity (not root)", 424 category=IdentityWarning) 425 return 426 427 user, group = str(user), str(group) 428 429 # resolve user 430 import pwd 431 try: 432 try: 433 userid = int(user) 434 except (TypeError, ValueError): 435 userid = pwd.getpwnam(user).pw_uid 436 else: 437 user = pwd.getpwuid(userid).pw_name 438 except KeyError, e: 439 raise IdentityError( 440 "User resolution problem of %r: %s" % (user, str(e)) 441 ) 442 443 # resolve group 444 import grp 445 try: 446 try: 447 groupid = int(group) 448 except (TypeError, ValueError): 449 groupid = grp.getgrnam(group).gr_gid 450 else: 451 group = grp.getgrgid(groupid).gr_name 452 except KeyError, e: 453 raise IdentityError( 454 "Group resolution problem of %r: %s" % (group, str(e)) 455 ) 456 457 # now do change our identity; group first as we might not have the 458 # permissions to do so after we left the power of root behind us. 459 _os.setgid(groupid) 460 try: 461 initgroups(user, groupid) 462 except NotImplementedError: 463 _warnings.warn("initgroups(3) is not implemented. You have to run " 464 "without supplemental groups or compile the wtf package " 465 "properly.", category=IdentityWarning) 466 _os.setuid(userid)
467 468
469 -def initgroups(username, gid):
470 """ 471 Implement initgroups(3) 472 473 :Parameters: 474 - `username`: The user name 475 - `gid`: The group id 476 477 :Types: 478 - `username`: ``str`` 479 - `gid`: ``int`` 480 481 :Exceptions: 482 - `OSError`: initgroups() didn't succeed 483 - `NotImplementedError`: initgroups is not implemented 484 (needs c-extension) 485 """ 486 # pylint: disable = W0613 487 488 raise NotImplementedError()
489 490 491 from wtf import c_override 492 cimpl = c_override('_wtf_cutil') 493 if cimpl is not None: 494 # pylint: disable = E1103 495 initgroups = cimpl.initgroups 496 del c_override, cimpl 497