From c1aad71dd3ba9acb7fc0e159816edb194752693c Mon Sep 17 00:00:00 2001
From: clement rouault <clement.rouault@exatrack.com>
Date: Mon, 31 Oct 2022 10:18:23 +0100
Subject: [PATCH] Improved integration of ws2_32 API + add demo_ws2_32.py
 sample

---
 .../definitions/defines/wsocket.txt           |  1 +
 samples/network/demo_ws2_32.py                | 36 +++++++++++++++++++
 windows/generated_def/__init__.py             |  2 ++
 windows/winproxy/apis/ws2_32.py               |  4 ++-
 4 files changed, 42 insertions(+), 1 deletion(-)
 create mode 100644 samples/network/demo_ws2_32.py

diff --git a/ctypes_generation/definitions/defines/wsocket.txt b/ctypes_generation/definitions/defines/wsocket.txt
index 6d95109..80381c9 100644
--- a/ctypes_generation/definitions/defines/wsocket.txt
+++ b/ctypes_generation/definitions/defines/wsocket.txt
@@ -7,6 +7,7 @@
 
 
 // INVALID_SOCKET  (SOCKET)(~0) // will change in 32 & 64 bits..
+// The final INVALID_SOCKET is setuped in __init__.py of generated def
 #define INVALID_SOCKET32   (0xffffffff)
 #define INVALID_SOCKET64   (0xffffffffffffffff)
 
diff --git a/samples/network/demo_ws2_32.py b/samples/network/demo_ws2_32.py
new file mode 100644
index 0000000..fa209d7
--- /dev/null
+++ b/samples/network/demo_ws2_32.py
@@ -0,0 +1,36 @@
+import windows.generated_def as gdef
+import windows.winproxy
+import ctypes
+
+
+WEBSITE = b"perdu.com"
+MESSAGE = b"""GET / HTTP/1.1\r\nHost: perdu.com\r\n\r\n"""
+
+x = gdef.WSADATA()
+windows.winproxy.WSAStartup(0x0202, x)
+
+print("=== WSADATA ===")
+windows.utils.sprint(x)
+
+hints = gdef.ADDRINFOA()
+hints.ai_family = gdef.AF_UNSPEC
+hints.ai_socktype = gdef.SOCK_STREAM
+hints.ai_protocol = gdef.IPPROTO_TCP
+
+result = gdef.PADDRINFOA()
+windows.winproxy.getaddrinfo(WEBSITE, b"80", hints, result);
+
+print("=== PADDRINFOA ===")
+windows.utils.sprint(result)
+
+
+connect_socket = windows.winproxy.socket(result[0].ai_family, result[0].ai_socktype, result[0].ai_protocol)
+res =  windows.winproxy.connect(connect_socket, result[0].ai_addr, result[0].ai_addrlen)
+windows.winproxy.send(connect_socket, MESSAGE)
+
+buf = ctypes.create_string_buffer(10000)
+windows.winproxy.recv(connect_socket, buf)
+print("Received:\n{0}".format(buf.value))
+
+windows.winproxy.closesocket(connect_socket)
+windows.winproxy.WSACleanup()
\ No newline at end of file
diff --git a/windows/generated_def/__init__.py b/windows/generated_def/__init__.py
index 3b49b2e..af1f9a2 100644
--- a/windows/generated_def/__init__.py
+++ b/windows/generated_def/__init__.py
@@ -29,6 +29,7 @@ if bitness() == 32:
     winstructs.IRQ_RESOURCE = winstructs.IRQ_RESOURCE_32
 
     # Socket
+    windef.WSADATA = winstructs.WSADATA32
     windef.INVALID_SOCKET = windef.INVALID_SOCKET32
 
 
@@ -53,6 +54,7 @@ else:
     winstructs.IRQ_RESOURCE = winstructs.IRQ_RESOURCE_64
 
     # Socket
+    windef.WSADATA = winstructs.WSADATA64
     windef.INVALID_SOCKET = windef.INVALID_SOCKET64
 
 from . import winfuncs
diff --git a/windows/winproxy/apis/ws2_32.py b/windows/winproxy/apis/ws2_32.py
index b4d8f63..e43fc47 100644
--- a/windows/winproxy/apis/ws2_32.py
+++ b/windows/winproxy/apis/ws2_32.py
@@ -18,7 +18,9 @@ def check_invalid_socket(func_name, result, func, args):
 
 @Ws2_32Proxy()
 def WSAStartup(wVersionRequested, lpWSAData):
-   return WSAStartup.ctypes_function(wVersionRequested, lpWSAData)
+    if isinstance(lpWSAData, gdef.WSADATA):
+        lpWSAData = ctypes.byref(lpWSAData) #  Not naturally done as lpWSAData is defined as a PVOID due to WSADATA32/WSADATA64 types
+    return WSAStartup.ctypes_function(wVersionRequested, lpWSAData)
 
 
 @Ws2_32Proxy()