Patch #934711: Expose platform-specific entropy.

loewis · loewis · commit dc3883f67189 · 2004-08-29T15:46:35.000Z
diff --git a/Doc/lib/libos.tex b/Doc/lib/libos.tex
@@ -1828,3 +1828,22 @@ \subsection{Miscellaneous System Information \label{os-path}}
 Also available via \module{os.path}.
 \versionadded{2.4}
 \end{datadesc}
+
+
+\subsection{Miscellaneous Functions \label{os-miscfunc}}
+
+\begin{funcdesc}{urandom}{n}
+Return a string of \var{n} random bytes suitable for cryptographic use.
+
+This function returns random bytes from an OS-specific 
+randomness source.  The returned data should be unpredictable enough for 
+cryptographic applications, though its exact quality depends on the OS 
+implementation.  On a UNIX-like system this will query /dev/urandom, and 
+on Windows it will use CryptGenRandom.  If a randomness source is not 
+found, \exception{NotImplementedError} will be raised.
+\versionadded{2.4}
+\end{funcdesc}
+
+
+
+
diff --git a/Lib/os.py b/Lib/os.py
@@ -656,3 +656,24 @@ def _pickle_statvfs_result(sr):
                      _make_statvfs_result)
 except NameError: # statvfs_result may not exist
     pass
+
+if not _exists("urandom"):
+    _urandomfd = None   
+    def urandom(n):
+        """urandom(n) -> str
+ 
+        Return a string of n random bytes suitable for cryptographic use.
+
+        """    
+        global _urandomfd
+        if not _urandomfd:
+            try:
+                _urandomfd = open("/dev/urandom", O_RDONLY)         
+            except:
+                _urandomfd = NotImplementedError
+        if _urandomfd is NotImplementedError:
+            raise NotImplementedError("/dev/urandom (or equivalent) not found")
+        bytes = ""
+        while len(bytes) < n:
+            bytes += read(_urandomfd, n - len(bytes))
+        return bytes
diff --git a/Lib/test/test_os.py b/Lib/test/test_os.py
@@ -343,6 +343,16 @@ def test_devnull(self):
         self.assertEqual(f.read(), '')
         f.close()
 
+class URandomTests (unittest.TestCase):
+    def test_urandom(self):
+        try:
+            self.assertEqual(len(os.urandom(1)), 1)
+            self.assertEqual(len(os.urandom(10)), 10)
+            self.assertEqual(len(os.urandom(100)), 100)
+            self.assertEqual(len(os.urandom(1000)), 1000)
+        except NotImplementedError:
+            pass
+
 def test_main():
     test_support.run_unittest(
         TemporaryFileTests,
@@ -351,6 +361,7 @@ def test_main():
         WalkTests,
         MakedirTests,
         DevNullTests,
+        URandomTests	
     )
 
 if __name__ == "__main__":
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -47,6 +47,9 @@ Core and builtins
 Extension modules
 -----------------
 
+- os.urandom has been added for systems that support sources of random
+  data.
+
 - Patch 1012740:  truncate() on a writeable cStringIO now resets the
   position to the end of the stream.  This is consistent with the original
   StringIO module and avoids inadvertently resurrecting data that was
diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c
@@ -228,7 +228,7 @@ extern int lstat(const char *, struct stat *);
 #include <io.h>
 #include <process.h>
 #include "osdefs.h"
-#define WIN32_LEAN_AND_MEAN
+#define _WIN32_WINNT 0x0400	  /* Needed for CryptoAPI on some systems */
 #include <windows.h>
 #include <shellapi.h>	/* for ShellExecute() */
 #define popen	_popen
@@ -7221,6 +7221,74 @@ posix_getloadavg(PyObject *self, PyObject *noargs)
 }
 #endif
 
+#ifdef MS_WINDOWS
+
+PyDoc_STRVAR(win32_urandom__doc__,
+"urandom(n) -> str\n\n\
+Return a string of n random bytes suitable for cryptographic use.");
+
+typedef BOOL (WINAPI *CRYPTACQUIRECONTEXTA)(HCRYPTPROV *phProv,\
+              LPCSTR pszContainer, LPCSTR pszProvider, DWORD dwProvType,\
+              DWORD dwFlags );
+typedef BOOL (WINAPI *CRYPTGENRANDOM)(HCRYPTPROV hProv, DWORD dwLen,\
+              BYTE *pbBuffer );
+
+static CRYPTGENRANDOM pCryptGenRandom = NULL;
+static HCRYPTPROV hCryptProv = 0;
+
+static PyObject* win32_urandom(PyObject *self, PyObject *args)
+{
+    int howMany = 0;
+    unsigned char* bytes = NULL;
+    PyObject* returnVal = NULL;
+
+    /* Read arguments */
+    if (!PyArg_ParseTuple(args, "i", &howMany))
+        return(NULL);
+
+    if (hCryptProv == 0) {
+        HINSTANCE hAdvAPI32 = NULL;
+        CRYPTACQUIRECONTEXTA pCryptAcquireContext = NULL;
+
+        /* Obtain handle to the DLL containing CryptoAPI
+           This should not fail	*/
+        if( (hAdvAPI32 = GetModuleHandle("advapi32.dll")) == NULL)
+            return win32_error("GetModuleHandle", NULL);
+
+        /* Obtain pointers to the CryptoAPI functions
+           This will fail on some early versions of Win95 */
+        pCryptAcquireContext=(CRYPTACQUIRECONTEXTA)GetProcAddress(hAdvAPI32,\
+                             "CryptAcquireContextA");
+        pCryptGenRandom=(CRYPTGENRANDOM)GetProcAddress(hAdvAPI32,\
+                        "CryptGenRandom");
+
+        if (pCryptAcquireContext == NULL || pCryptGenRandom == NULL)
+            return PyErr_Format(PyExc_NotImplementedError,\
+                                "CryptGenRandom not found");  
+
+	    /* Acquire context */
+        if(!pCryptAcquireContext(&hCryptProv, NULL, NULL, PROV_RSA_FULL,
+                                 CRYPT_VERIFYCONTEXT))
+            return win32_error("CryptAcquireContext", NULL);
+    }
+
+    /* Allocate bytes */
+    if ((bytes = (unsigned char*)PyMem_Malloc(howMany)) == NULL)
+        return PyErr_NoMemory();
+
+    /* Get random data */
+    if (!pCryptGenRandom(hCryptProv, howMany, bytes)) {
+        PyMem_Free(bytes);
+        return win32_error("CryptGenRandom", NULL);
+    }
+
+    /* Build return value */
+    returnVal = PyString_FromStringAndSize(bytes, howMany);
+    PyMem_Free(bytes);
+
+    return returnVal;
+}
+#endif
 
 static PyMethodDef posix_methods[] = {
 	{"access",	posix_access, METH_VARARGS, posix_access__doc__},
@@ -7506,6 +7574,9 @@ static PyMethodDef posix_methods[] = {
 #ifdef HAVE_GETLOADAVG
 	{"getloadavg",	posix_getloadavg, METH_NOARGS, posix_getloadavg__doc__},
 #endif
+ #ifdef MS_WINDOWS
+ 	{"urandom", win32_urandom, METH_VARARGS, win32_urandom__doc__},
+ #endif
 	{NULL,		NULL}		 /* Sentinel */
 };
 
