From 01dca258ee480ea085d2c80a0c8711acb48b6634 Mon Sep 17 00:00:00 2001
From: Jondy Zhao <jondy.zhao@gmail.com>
Date: Fri, 17 May 2013 17:17:07 +0800
Subject: [PATCH] Add method: userinfo, usagereport

---
 windows/netreport/README       |  10 ++
 windows/netreport/setup.py     |   5 +-
 windows/netreport/src/netuse.c | 256 +++++++++++++++++++++++++++++++--
 3 files changed, 259 insertions(+), 12 deletions(-)

diff --git a/windows/netreport/README b/windows/netreport/README
index 24ab580..7aff349 100644
--- a/windows/netreport/README
+++ b/windows/netreport/README
@@ -1 +1,11 @@
 Net Resource Usage Report For Windows
+
+Build extentions
+================
+
+$ python setup.py build
+
+Test basic functions
+====================
+
+$ python tests.py
diff --git a/windows/netreport/setup.py b/windows/netreport/setup.py
index cd81341..a629d54 100644
--- a/windows/netreport/setup.py
+++ b/windows/netreport/setup.py
@@ -30,9 +30,8 @@ if sys.platform.startswith("cygwin"):
                             sources=['src/netuse.c'],
                             define_macros=[('_WIN32_WINNT', '0x0503'),
                                            ('USE_SYS_TYPES_FD_SET', 1)],
-                            libraries=["psapi", "kernel32", "advapi32",
-                                       "shell32", "netapi32", "iphlpapi",
-                                       "wtsapi32"],
+                            libraries=["kernel32", "advapi32", "shell32",
+                                       "netapi32", "mpr"],
                             #extra_compile_args=["/Z7"],
                             #extra_link_args=["/DEBUG"]
                             )]
diff --git a/windows/netreport/src/netuse.c b/windows/netreport/src/netuse.c
index de8ed88..eee33e1 100644
--- a/windows/netreport/src/netuse.c
+++ b/windows/netreport/src/netuse.c
@@ -25,35 +25,273 @@
 #include <string.h>
 #include <sys/cygwin.h>
 
-/* Avoid select function conflict in the winsock2.h */
-#define __INSIDE_CYGWIN__
 #include <windows.h>
+#include <lm.h>
+#include <winnetwk.h>
 
 #define MALLOC(x) HeapAlloc(GetProcessHeap(), 0, (x))
 #define FREE(x) HeapFree(GetProcessHeap(), 0, (x))
 
+#define MAX_USERBUFFER_SIZE 1024
+
+static char userinfo[MAX_USERBUFFER_SIZE] = { 0 };
+static char * logonuser = NULL;
+static char * logondomain = NULL;
+static char * logonserver = NULL;
+
+static size_t
+wchar2mchar(wchar_t *ws, char *buffer, size_t size)
+{
+  size_t len;
+  len = WideCharToMultiByte(CP_ACP,
+                            0,
+                            ws,
+                            -1,
+                            NULL,
+                            0,
+                            NULL,
+                            NULL
+                            );
+  if (len + 1 > size)
+    return -1;
+  if (WideCharToMultiByte(CP_ACP,
+                          0,
+                          ws,
+                          -1,
+                          buffer,
+                          len,
+                          NULL,
+                          NULL
+                          ) == 0)
+    return -1;
+  return len + 1;
+}
+
+static PyObject *
+netuse_user_info(PyObject *self, PyObject *args)
+{
+   DWORD dwLevel = 1;
+   LPWKSTA_USER_INFO_1 pBuf = NULL;
+   NET_API_STATUS nStatus;
+   //
+   // Call the NetWkstaUserGetInfo function;
+   //  specify level 1.
+   //
+   nStatus = NetWkstaUserGetInfo(NULL,
+                                 dwLevel,
+                                 (LPBYTE *)&pBuf);
+   //
+   // If the call succeeds, print the information
+   //  about the logged-on user.
+   //
+   if (nStatus == NERR_Success) {
+     if (pBuf != NULL) {
+       size_t size = MAX_USERBUFFER_SIZE;
+       size_t len;
+       logonuser = userinfo;
+       len = wchar2mchar(pBuf->wkui1_username, logonuser, size);
+       if (len == -1) {
+         PyErr_SetString(PyExc_RuntimeError, "Unicode convertion error");
+         return NULL;
+       }
+       size -= len;
+       logondomain = logonuser + len;
+       len = wchar2mchar(pBuf->wkui1_logon_domain, logondomain, size);
+       if (len == -1) {
+         PyErr_SetString(PyExc_RuntimeError, "Unicode convertion error");
+         return NULL;
+       }
+       size -= len;
+       logonserver = logondomain + len;
+       len = wchar2mchar(pBuf->wkui1_logon_server, logonserver, size);
+       if (len == -1) {
+         PyErr_SetString(PyExc_RuntimeError, "Unicode convertion error");
+         return NULL;
+       }
+     }
+   }
+   // Otherwise, print the system error.
+   //
+   else {
+     PyErr_Format(PyExc_RuntimeError,
+                  "A system error has occurred: %ld",
+                  nStatus
+                  );
+     return NULL;
+   }
+   //
+   // Free the allocated memory.
+   //
+   if (pBuf != NULL) {
+     NetApiBufferFree(pBuf);
+     return Py_BuildValue("sss", logonuser, logondomain, logonserver);
+   }
+
+   PyErr_SetString(PyExc_RuntimeError, "No logon user information");
+   return NULL;
+}
+
 static PyObject *
-netuse_init(PyObject *self, PyObject *args)
+netuse_map_drive(PyObject *self, PyObject *args)
 {
   return NULL;
 }
 
+static PyObject *
+netuse_usage_report(PyObject *self, PyObject *args)
+{
+  char * servername = NULL;
+  PyObject *retvalue = NULL;
+  DWORD bitmasks;
+  char chdrive = '@';
+  char  drivepath[4] = { 'A', ':', '\\', 0 };
+  char  drivename[3] = { 'A', ':', 0 };
+  ULARGE_INTEGER lFreeBytesAvailable;
+  ULARGE_INTEGER lTotalNumberOfBytes;
+  ULARGE_INTEGER lTotalNumberOfFreeBytes;
+
+  char szRemoteName[MAX_PATH];
+  DWORD dwResult, cchBuff = MAX_PATH;
+  DWORD serverlen = 0;
+
+  if (! PyArg_ParseTuple(args, "|s", &servername)) {
+    return NULL;
+  }
+
+  if (servername) 
+    serverlen = strlen(servername);
+
+  bitmasks = GetLogicalDrives();
+  if (bitmasks == 0) {
+     PyErr_Format(PyExc_RuntimeError,
+                  "A system error has occurred in GetLogicalDrives: %ld",
+                  GetLastError()
+                  );
+     return NULL;
+  }
+
+  retvalue = PyList_New(0);
+  if (retvalue == NULL)
+    return NULL;
+
+  while (bitmasks) {
+    ++ chdrive;
+    drivepath[0] = chdrive;
+    drivename[0] = chdrive;
+
+    if ((bitmasks & 1L) == 0) {
+      bitmasks >>= 1;
+      continue;
+    }
+
+    bitmasks >>= 1;
+    switch (GetDriveType(drivepath)) {
+      case DRIVE_FIXED:
+      case DRIVE_REMOTE:
+        break;
+      default:
+        continue;
+      }
+
+    dwResult = WNetGetConnection(drivename,
+                                 szRemoteName,
+                                 &cchBuff
+                                 );
+    if (dwResult == NO_ERROR) {
+      if (servername) {
+        if ((cchBuff < serverlen + 3) ||
+            (strncmp(servername, szRemoteName+2, serverlen) != 0) ||
+            (szRemoteName[serverlen + 2] != '\\')
+            )
+        continue;
+      }
+    }
+
+    // The device is not currently connected, but it is a persistent connection.
+    else if (dwResult == ERROR_CONNECTION_UNAVAIL) {
+      continue;
+    }
+
+    // The device is not a redirected device.
+    else if (dwResult == ERROR_NOT_CONNECTED) {
+      continue;
+    }
+
+    else {
+      PyErr_Format(PyExc_RuntimeError,
+                   "A system error has occurred in WNetGetConnection: %ld",
+                   GetLastError()
+                   );
+      Py_XDECREF(retvalue);
+      return NULL;
+    }
+    
+    if (GetDiskFreeSpaceEx(drivepath,
+                           &lFreeBytesAvailable,
+                           &lTotalNumberOfBytes,
+                           &lTotalNumberOfFreeBytes
+                           )) {
+      PyObject *pobj = Py_BuildValue("ssLLL",
+                                     drivename,
+                                     szRemoteName,
+                                     lFreeBytesAvailable,
+                                     lTotalNumberOfFreeBytes,
+                                     lTotalNumberOfBytes
+                                     );
+      if (PyList_Append(retvalue, pobj) == -1) {
+        Py_XDECREF(retvalue);
+        return NULL;
+      }
+    }
+    else {      
+     PyErr_Format(PyExc_RuntimeError,
+                  "A system error has occurred in GetDiskFreeSpaceEx(%s): %ld",
+                  drivepath,
+                  GetLastError()
+                  );
+     Py_XDECREF(retvalue);
+     return NULL;
+    }
+  }
+
+  return retvalue;
+}
+
 static PyMethodDef NetUseMethods[] = {
   {
-    "init",
-    netuse_init,
+    "userinfo",
+    netuse_user_info,
+    METH_VARARGS,
+    (
+     "userinfo()\n\n"
+     "Get the logon user information, return a tuple:\n"
+     "(user, domain, server).\n"
+     )
+  },
+  {
+    "mapdrive",
+    netuse_map_drive,
+    METH_VARARGS,
+    (
+     "mapdrive()\n\n"
+     "Create mapped drive from server shared folder\n"
+     )
+  },
+  {
+    "usagereport",
+    netuse_usage_report,
     METH_VARARGS,
     (
-     "init()\n\n"
-     "Initialize an inotify instance and return a PyCObject, When this\n"
-     "PyCObject is reclaimed, GC will free the memory.\n"
+     "usagereport()\n\n"
+     "Return a tuple to report all the mapped drive information:\n"
+     "(user, domain, drive, usage, total).\n"
      )
   },
   {NULL, NULL, 0, NULL}
 };
 
 
-PyMODINIT_FUNC netuse(void)
+PyMODINIT_FUNC initnetuse(void)
 {
   PyObject* module;
   module = Py_InitModule3("netuse",
-- 
2.30.9