openerp/tools/osutil.py

   1 # -*- coding: utf-8 -*-
   2 ##############################################################################
   3 #
   4 #    OpenERP, Open Source Management Solution
   5 #    Copyright (C) 2004-2009 Tiny SPRL (<http://tiny.be>).
   6 #
   7 #    This program is free software: you can redistribute it and/or modify
   8 #    it under the terms of the GNU Affero General Public License as
   9 #    published by the Free Software Foundation, either version 3 of the
  10 #    License, or (at your option) any later version.
  11 #
  12 #    This program is distributed in the hope that it will be useful,
  13 #    but WITHOUT ANY WARRANTY; without even the implied warranty of
  14 #    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  15 #    GNU Affero General Public License for more details.
  16 #
  17 #    You should have received a copy of the GNU Affero General Public License
  18 #    along with this program.  If not, see <http://www.gnu.org/licenses/>.
  19 #
  20 ##############################################################################
  21
  22 """
  23 Some functions related to the os and os.path module
  24 """
  25
  26 import os
  27 from os.path import join as opj
  28
  29 if os.name == 'nt':
  30     import ctypes
  31     import win32service as ws
  32     import win32serviceutil as wsu
  33
  34
  35 def listdir(dir, recursive=False):
  36     """Allow to recursively get the file listing"""
  37     dir = os.path.normpath(dir)
  38     if not recursive:
  39         return os.listdir(dir)
  40
  41     res = []
  42     for root, dirs, files in walksymlinks(dir):
  43         root = root[len(dir)+1:]
  44         res.extend([opj(root, f) for f in files])
  45     return res
  46
  47 def walksymlinks(top, topdown=True, onerror=None):
  48     """
  49     same as os.walk but follow symlinks
  50     attention: all symlinks are walked before all normals directories
  51     """
  52     for dirpath, dirnames, filenames in os.walk(top, topdown, onerror):
  53         if topdown:
  54             yield dirpath, dirnames, filenames
  55
  56         symlinks = filter(lambda dirname: os.path.islink(os.path.join(dirpath, dirname)), dirnames)
  57         for s in symlinks:
  58             for x in walksymlinks(os.path.join(dirpath, s), topdown, onerror):
  59                 yield x
  60
  61         if not topdown:
  62             yield dirpath, dirnames, filenames
  63
  64
  65 if os.name != 'nt':
  66     getppid = os.getppid
  67     is_running_as_nt_service = lambda: False
  68 else:
  69     # based on http://mail.python.org/pipermail/python-win32/2007-June/006174.html
  70     _TH32CS_SNAPPROCESS = 0x00000002
  71     class _PROCESSENTRY32(ctypes.Structure):
  72         _fields_ = [("dwSize", ctypes.c_ulong),
  73                     ("cntUsage", ctypes.c_ulong),
  74                     ("th32ProcessID", ctypes.c_ulong),
  75                     ("th32DefaultHeapID", ctypes.c_ulong),
  76                     ("th32ModuleID", ctypes.c_ulong),
  77                     ("cntThreads", ctypes.c_ulong),
  78                     ("th32ParentProcessID", ctypes.c_ulong),
  79                     ("pcPriClassBase", ctypes.c_ulong),
  80                     ("dwFlags", ctypes.c_ulong),
  81                     ("szExeFile", ctypes.c_char * 260)]
  82
  83     def getppid():
  84         CreateToolhelp32Snapshot = ctypes.windll.kernel32.CreateToolhelp32Snapshot
  85         Process32First = ctypes.windll.kernel32.Process32First
  86         Process32Next = ctypes.windll.kernel32.Process32Next
  87         CloseHandle = ctypes.windll.kernel32.CloseHandle
  88         hProcessSnap = CreateToolhelp32Snapshot(_TH32CS_SNAPPROCESS, 0)
  89         current_pid = os.getpid()
  90         try:
  91             pe32 = _PROCESSENTRY32()
  92             pe32.dwSize = ctypes.sizeof(_PROCESSENTRY32)
  93             if not Process32First(hProcessSnap, ctypes.byref(pe32)):
  94                 raise OSError('Failed getting first process.')
  95             while True:
  96                 if pe32.th32ProcessID == current_pid:
  97                     return pe32.th32ParentProcessID
  98                 if not Process32Next(hProcessSnap, ctypes.byref(pe32)):
  99                     return None
 100         finally:
 101             CloseHandle(hProcessSnap)
 102
 103     from contextlib import contextmanager
 104     from openerp.release import nt_service_name
 105
 106     def is_running_as_nt_service():
 107         @contextmanager
 108         def close_srv(srv):
 109             try:
 110                 yield srv
 111             finally:
 112                 ws.CloseServiceHandle(srv)
 113
 114         with close_srv(ws.OpenSCManager(None, None, ws.SC_MANAGER_ALL_ACCESS)) as hscm:
 115             with close_srv(wsu.SmartOpenService(hscm, nt_service_name, ws.SERVICE_ALL_ACCESS)) as hs:
 116                 info = ws.QueryServiceStatusEx(hs)
 117                 return info['ProcessId'] == getppid()
 118
 119 if __name__ == '__main__':
 120     from pprint import pprint as pp
 121     pp(listdir('../report', True))
 122
 123 # vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4: