* [KVM-AUTOTEST PATCH v4] [RFC] KVM test: rss.cpp: add file transfer support
@ 2010-07-04 13:42 Michael Goldish
2010-07-04 13:42 ` [KVM-AUTOTEST PATCH v4] [RFC] KVM test: add python client for rss file transfer services Michael Goldish
0 siblings, 1 reply; 3+ messages in thread
From: Michael Goldish @ 2010-07-04 13:42 UTC (permalink / raw)
To: autotest, kvm; +Cc: Michael Goldish
Enable RSS to send/receive files and directory trees (recursively).
See protocol details in rss.cpp.
Changes from v3:
- Protocol change: instead of sending a file in one big packet, send it in
multiple chunks. The last chunk of a file is identified by being smaller.
This solves a problem with transfers of files that happen to be growing or
shrinking while being transferred.
- Change the way messages are displayed in the text box: instead of adding them
immediately and redrawing the text box for each message, queue the messages
in a buffer and flush it every 250ms or when it's full. This makes it
possible to log everything without using a significant amount of CPU time,
even for transfers of thousands of small files.
- Change text box limit to 262144 chars.
- Allow resizing the main window.
- Reuse some code by putting it in an Accept() function.
Changes from v2:
- Use ports 10022 and 10023 by default instead of 22 and 23.
Changes from v1:
- Expand environment variables (e.g. %WinDir%) in all paths.
- Change text box limit to 16384 chars.
- When text box is full, clear 3/4 of old text instead of 1/2 (looks prettier).
- Use const char * instead of char * where appropriate.
Signed-off-by: Michael Goldish <mgoldish@redhat.com>
---
| 1449 ++++++++++++++++++++++++++++-------------
1 files changed, 990 insertions(+), 459 deletions(-)
--git a/client/tests/kvm/deps/rss.cpp b/client/tests/kvm/deps/rss.cpp
index 66d9a5b..26c5ed6 100644
--- a/client/tests/kvm/deps/rss.cpp
+++ b/client/tests/kvm/deps/rss.cpp
@@ -1,459 +1,990 @@
-// Simple remote shell server
-// Author: Michael Goldish <mgoldish@redhat.com>
-// Much of the code here was adapted from Microsoft code samples.
-
-// Usage: rss.exe [port]
-// If no port is specified the default is 22.
-
-#define _WIN32_WINNT 0x0500
-
-#include <windows.h>
-#include <winsock2.h>
-#include <stdio.h>
-
-#pragma comment(lib, "ws2_32.lib")
-
-int port = 22;
-
-HWND hMainWindow = NULL;
-HWND hTextBox = NULL;
-
-struct client_info {
- SOCKET socket;
- sockaddr_in addr;
- int pid;
- HWND hwnd;
- HANDLE hJob;
- HANDLE hChildOutputRead;
- HANDLE hThreadChildToSocket;
-};
-
-void ExitOnError(char *message, BOOL winsock = 0)
-{
- LPVOID system_message;
- char buffer[512];
-
- int error_code;
- if (winsock)
- error_code = WSAGetLastError();
- else
- error_code = GetLastError();
-
- WSACleanup();
-
- FormatMessage(
- FORMAT_MESSAGE_ALLOCATE_BUFFER|FORMAT_MESSAGE_FROM_SYSTEM,
- NULL, error_code, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
- (LPTSTR)&system_message, 0, NULL);
-
- sprintf(buffer,
- "%s!\n"
- "Error code = %d\n"
- "Error message = %s",
- message, error_code, (char *)system_message);
-
- MessageBox(hMainWindow, buffer, "Error", MB_OK | MB_ICONERROR);
-
- LocalFree(system_message);
- ExitProcess(1);
-}
-
-void AppendMessage(char *message)
-{
- int length = GetWindowTextLength(hTextBox);
- SendMessage(hTextBox, EM_SETSEL, (WPARAM)length, (LPARAM)length);
- SendMessage(hTextBox, EM_REPLACESEL, (WPARAM)FALSE, (LPARAM)message);
-}
-
-void FormatStringForPrinting(char *dst, char *src, int size)
-{
- int j = 0;
-
- for (int i = 0; i < size && src[i]; i++) {
- if (src[i] == '\n') {
- dst[j++] = '\\';
- dst[j++] = 'n';
- } else if (src[i] == '\r') {
- dst[j++] = '\\';
- dst[j++] = 'r';
- } else if (src[i] == '\t') {
- dst[j++] = '\\';
- dst[j++] = 't';
- } else if (src[i] == '\\') {
- dst[j++] = '\\';
- dst[j++] = '\\';
- } else dst[j++] = src[i];
- }
- dst[j] = 0;
-}
-
-char* GetClientIPAddress(client_info *ci)
-{
- char *address = inet_ntoa(ci->addr.sin_addr);
- if (address)
- return address;
- else
- return "unknown";
-}
-
-DWORD WINAPI ChildToSocket(LPVOID client_info_ptr)
-{
- char buffer[1024], message[1024];
- client_info ci;
- DWORD bytes_read;
- int bytes_sent;
-
- memcpy(&ci, client_info_ptr, sizeof(ci));
-
- while (1) {
- // Read data from the child's STDOUT/STDERR pipes
- if (!ReadFile(ci.hChildOutputRead,
- buffer, sizeof(buffer),
- &bytes_read, NULL) || !bytes_read) {
- if (GetLastError() == ERROR_BROKEN_PIPE)
- break; // Pipe done -- normal exit path
- else
- ExitOnError("ReadFile failed"); // Something bad happened
- }
- // Send data to the client
- bytes_sent = send(ci.socket, buffer, bytes_read, 0);
- /*
- // Make sure all the data was sent
- if (bytes_sent != bytes_read) {
- sprintf(message,
- "ChildToSocket: bytes read (%d) != bytes sent (%d)",
- bytes_read, bytes_sent);
- ExitOnError(message, 1);
- }
- */
- }
-
- AppendMessage("Child exited\r\n");
- shutdown(ci.socket, SD_BOTH);
-
- return 0;
-}
-
-DWORD WINAPI SocketToChild(LPVOID client_info_ptr)
-{
- char buffer[256], formatted_buffer[768];
- char message[1024], client_info_str[256];
- client_info ci;
- DWORD bytes_written;
- int bytes_received;
-
- memcpy(&ci, client_info_ptr, sizeof(ci));
-
- sprintf(client_info_str, "address %s, port %d",
- GetClientIPAddress(&ci), ci.addr.sin_port);
-
- sprintf(message, "New client connected (%s)\r\n", client_info_str);
- AppendMessage(message);
-
- while (1) {
- // Receive data from the socket
- ZeroMemory(buffer, sizeof(buffer));
- bytes_received = recv(ci.socket, buffer, sizeof(buffer), 0);
- if (bytes_received <= 0)
- break;
- // Report the data received
- FormatStringForPrinting(formatted_buffer, buffer, sizeof(buffer));
- sprintf(message, "Client (%s) entered text: \"%s\"\r\n",
- client_info_str, formatted_buffer);
- AppendMessage(message);
- // Send the data as a series of WM_CHAR messages to the console window
- for (int i=0; i<bytes_received; i++) {
- SendMessage(ci.hwnd, WM_CHAR, (WPARAM)buffer[i], 0);
- SendMessage(ci.hwnd, WM_SETFOCUS, 0, 0);
- }
- }
-
- sprintf(message, "Client disconnected (%s)\r\n", client_info_str);
- AppendMessage(message);
-
- // Attempt to terminate the child's process tree:
- // Using taskkill (where available)
- sprintf(buffer, "taskkill /PID %d /T /F", ci.pid);
- system(buffer);
- // .. and using TerminateJobObject()
- TerminateJobObject(ci.hJob, 0);
- // Wait for the ChildToSocket thread to terminate
- WaitForSingleObject(ci.hThreadChildToSocket, 10000);
- // In case the thread refuses to exit -- terminate it
- TerminateThread(ci.hThreadChildToSocket, 0);
- // Close the socket
- shutdown(ci.socket, SD_BOTH);
- closesocket(ci.socket);
-
- // Close unnecessary handles
- CloseHandle(ci.hJob);
- CloseHandle(ci.hThreadChildToSocket);
- CloseHandle(ci.hChildOutputRead);
-
- AppendMessage("SocketToChild thread exited\r\n");
-
- return 0;
-}
-
-void PrepAndLaunchRedirectedChild(client_info *ci,
- HANDLE hChildStdOut,
- HANDLE hChildStdErr)
-{
- PROCESS_INFORMATION pi;
- STARTUPINFO si;
-
- // Allocate a new console for the child
- HWND hwnd = GetForegroundWindow();
- FreeConsole();
- AllocConsole();
- ShowWindow(GetConsoleWindow(), SW_HIDE);
- if (hwnd)
- SetForegroundWindow(hwnd);
-
- // Set up the start up info struct.
- ZeroMemory(&si, sizeof(STARTUPINFO));
- si.cb = sizeof(STARTUPINFO);
- si.dwFlags = STARTF_USESTDHANDLES | STARTF_USESHOWWINDOW;
- si.hStdOutput = hChildStdOut;
- si.hStdInput = GetStdHandle(STD_INPUT_HANDLE);
- si.hStdError = hChildStdErr;
- // Use this if you want to hide the child:
- si.wShowWindow = SW_HIDE;
- // Note that dwFlags must include STARTF_USESHOWWINDOW if you want to
- // use the wShowWindow flags.
-
- // Launch the process that you want to redirect.
- if (!CreateProcess(NULL, "cmd.exe", NULL, NULL, TRUE,
- 0, NULL, "C:\\", &si, &pi))
- ExitOnError("CreateProcess failed");
-
- // Close any unnecessary handles.
- if (!CloseHandle(pi.hThread))
- ExitOnError("CloseHandle failed");
-
- // Keep the process ID
- ci->pid = pi.dwProcessId;
- // Assign the process to a newly created JobObject
- ci->hJob = CreateJobObject(NULL, NULL);
- AssignProcessToJobObject(ci->hJob, pi.hProcess);
- // Keep the console window's handle
- ci->hwnd = GetConsoleWindow();
-
- // Detach from the child's console
- FreeConsole();
-}
-
-void SpawnSession(client_info *ci)
-{
- HANDLE hOutputReadTmp, hOutputRead, hOutputWrite;
- HANDLE hErrorWrite;
- SECURITY_ATTRIBUTES sa;
-
- // Set up the security attributes struct.
- sa.nLength = sizeof(SECURITY_ATTRIBUTES);
- sa.lpSecurityDescriptor = NULL;
- sa.bInheritHandle = TRUE;
-
- // Create the child output pipe.
- if (!CreatePipe(&hOutputReadTmp, &hOutputWrite, &sa, 0))
- ExitOnError("CreatePipe failed");
-
- // Create a duplicate of the output write handle for the std error
- // write handle. This is necessary in case the child application
- // closes one of its std output handles.
- if (!DuplicateHandle(GetCurrentProcess(), hOutputWrite,
- GetCurrentProcess(), &hErrorWrite, 0,
- TRUE, DUPLICATE_SAME_ACCESS))
- ExitOnError("DuplicateHandle failed");
-
- // Create new output read handle and the input write handles. Set
- // the Properties to FALSE. Otherwise, the child inherits the
- // properties and, as a result, non-closeable handles to the pipes
- // are created.
- if (!DuplicateHandle(GetCurrentProcess(), hOutputReadTmp,
- GetCurrentProcess(),
- &hOutputRead, // Address of new handle.
- 0, FALSE, // Make it uninheritable.
- DUPLICATE_SAME_ACCESS))
- ExitOnError("DuplicateHandle failed");
-
- // Close inheritable copies of the handles you do not want to be
- // inherited.
- if (!CloseHandle(hOutputReadTmp))
- ExitOnError("CloseHandle failed");
-
- PrepAndLaunchRedirectedChild(ci, hOutputWrite, hErrorWrite);
-
- ci->hChildOutputRead = hOutputRead;
-
- // Close pipe handles (do not continue to modify the parent).
- // You need to make sure that no handles to the write end of the
- // output pipe are maintained in this process or else the pipe will
- // not close when the child process exits and the ReadFile will hang.
- if (!CloseHandle(hOutputWrite)) ExitOnError("CloseHandle failed");
- if (!CloseHandle(hErrorWrite)) ExitOnError("CloseHandle failed");
-}
-
-DWORD WINAPI ListenThread(LPVOID param)
-{
- WSADATA wsaData;
- SOCKET ListenSocket = INVALID_SOCKET;
- sockaddr_in addr;
- int result, addrlen;
- client_info ci;
- HANDLE hThread;
-
- // Initialize Winsock
- result = WSAStartup(MAKEWORD(2,2), &wsaData);
- if (result)
- ExitOnError("Winsock initialization failed");
-
- // Create socket
- ListenSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
- if (ListenSocket == INVALID_SOCKET)
- ExitOnError("Socket creation failed", 1);
-
- // Bind the socket
- addr.sin_family = AF_INET;
- addr.sin_addr.s_addr = htonl(INADDR_ANY);
- addr.sin_port = htons(port);
-
- result = bind(ListenSocket, (sockaddr *)&addr, sizeof(addr));
- if (result == SOCKET_ERROR)
- ExitOnError("bind failed", 1);
-
- // Start listening for incoming connections
- result = listen(ListenSocket, SOMAXCONN);
- if (result == SOCKET_ERROR)
- ExitOnError("listen failed", 1);
-
- // Inform the user
- AppendMessage("Waiting for clients to connect...\r\n");
-
- while (1) {
- addrlen = sizeof(ci.addr);
- ci.socket = accept(ListenSocket, (sockaddr *)&ci.addr, &addrlen);
- if (ci.socket == INVALID_SOCKET) {
- if (WSAGetLastError() == WSAEINTR)
- break;
- else
- ExitOnError("accept failed", 1);
- }
-
- // Under heavy load, spawning cmd.exe might take a while, so tell the
- // client to be patient
- char *message = "Please wait...\r\n";
- send(ci.socket, message, strlen(message), 0);
- // Spawn a new redirected cmd.exe process
- SpawnSession(&ci);
- // Start transferring data from the child process to the client
- hThread = CreateThread(NULL, 0, ChildToSocket, (LPVOID)&ci, 0, NULL);
- if (!hThread)
- ExitOnError("Could not create ChildToSocket thread");
- ci.hThreadChildToSocket = hThread;
- // ... and from the client to the child process
- hThread = CreateThread(NULL, 0, SocketToChild, (LPVOID)&ci, 0, NULL);
- if (!hThread)
- ExitOnError("Could not create SocketToChild thread");
- }
-
- return 0;
-}
-
-LRESULT CALLBACK WndProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam)
-{
- RECT rect;
- HANDLE hListenThread;
-
- switch (msg) {
- case WM_CREATE:
- // Create text box
- GetClientRect(hwnd, &rect);
- hTextBox = CreateWindowEx(WS_EX_CLIENTEDGE,
- "EDIT", "",
- WS_CHILD|WS_VISIBLE|WS_VSCROLL|
- ES_MULTILINE|ES_AUTOVSCROLL,
- 20, 20,
- rect.right - 40,
- rect.bottom - 40,
- hwnd,
- NULL,
- GetModuleHandle(NULL),
- NULL);
- if (!hTextBox)
- ExitOnError("Could not create text box");
-
- // Set the font
- SendMessage(hTextBox, WM_SETFONT,
- (WPARAM)GetStockObject(DEFAULT_GUI_FONT),
- MAKELPARAM(FALSE, 0));
-
- // Start the listening thread
- hListenThread =
- CreateThread(NULL, 0, ListenThread, NULL, 0, NULL);
- if (!hListenThread)
- ExitOnError("Could not create server thread");
- break;
-
- case WM_DESTROY:
- WSACleanup();
- PostQuitMessage(0);
- break;
-
- default:
- return DefWindowProc(hwnd, msg, wParam, lParam);
- }
-
- return 0;
-}
-
-int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance,
- LPSTR lpCmdLine, int nShowCmd)
-{
- WNDCLASSEX wc;
- MSG msg;
-
- if (strlen(lpCmdLine))
- sscanf(lpCmdLine, "%d", &port);
-
- // Make sure the firewall is disabled
- system("netsh firewall set opmode disable");
-
- // Create the window class
- wc.cbSize = sizeof(WNDCLASSEX);
- wc.style = CS_HREDRAW | CS_VREDRAW;
- wc.lpfnWndProc = WndProc;
- wc.cbClsExtra = 0;
- wc.cbWndExtra = 0;
- wc.hInstance = hInstance;
- wc.hIcon = LoadIcon(NULL, IDI_APPLICATION);
- wc.hIconSm = LoadIcon(NULL, IDI_APPLICATION);
- wc.hbrBackground = (HBRUSH)(COLOR_BTNFACE + 1);
- wc.lpszMenuName = NULL;
- wc.lpszClassName = "RemoteShellServerWindowClass";
- wc.hCursor = LoadCursor(NULL, IDC_ARROW);
-
- if (!RegisterClassEx(&wc))
- ExitOnError("Could not register window class");
-
- // Create the main window
- hMainWindow =
- CreateWindow("RemoteShellServerWindowClass",
- "Remote Shell Server",
- WS_OVERLAPPED|WS_CAPTION|WS_SYSMENU|WS_MINIMIZEBOX,
- 20, 20, 500, 300,
- NULL, NULL, hInstance, NULL);
- if (!hMainWindow)
- ExitOnError("Could not create window");
-
- ShowWindow(hMainWindow, SW_SHOWMINNOACTIVE);
- UpdateWindow(hMainWindow);
-
- // Main message loop
- while (GetMessage(&msg, NULL, 0, 0)) {
- TranslateMessage(&msg);
- DispatchMessage(&msg);
- }
-
- ExitProcess(0);
-}
+// Simple remote shell server (and file transfer server)
+// Author: Michael Goldish <mgoldish@redhat.com>
+// Much of the code here was adapted from Microsoft code samples.
+
+// Usage: rss.exe [shell port] [file transfer port]
+// If no shell port is specified the default is 10022.
+// If no file transfer port is specified the default is 10023.
+
+// Definitions:
+// A 'msg' is a 32 bit integer.
+// A 'packet' is a 32 bit unsigned integer followed by a string of bytes.
+// The 32 bit integer indicates the length of the string.
+
+// Protocol for file transfers:
+//
+// When uploading files/directories to the server:
+// 1. The client connects.
+// 2. The server sends RSS_MAGIC.
+// 3. The client sends the chunk size for file transfers (a 32 bit integer
+// between 512 and 1048576 indicating the size in bytes).
+// 4. The client sends RSS_SET_PATH, followed by a packet (as defined above)
+// containing the path (in the server's filesystem) where files and/or
+// directories are to be stored.
+// Uploading a file (optional, can be repeated many times):
+// 5. The client sends RSS_CREATE_FILE, followed by a packet containing the
+// filename (filename only, without a path), followed by a series of
+// packets (called chunks) containing the file's contents. The size of
+// each chunk is the size set by the client in step 3, except for the
+// last chunk, which must be smaller.
+// Uploading a directory (optional, can be repeated many times):
+// 6. The client sends RSS_CREATE_DIR, followed by a packet containing the
+// name of the directory to be created (directory name only, without a
+// path).
+// 7. The client uploads files and directories to the new directory (using
+// steps 5, 6, 8).
+// 8. The client sends RSS_LEAVE_DIR.
+// 9. The client sends RSS_DONE and waits for a response.
+// 10. The server sends RSS_OK to indicate that it's still listening.
+// 11. Steps 4-10 are repeated as many times as necessary.
+// 12. The client disconnects.
+// If a critical error occurs at any time, the server may send RSS_ERROR
+// followed by a packet containing an error message, and the connection is
+// closed.
+//
+// When downloading files from the server:
+// 1. The client connects.
+// 2. The server sends RSS_MAGIC.
+// 3. The client sends the chunk size for file transfers (a 32 bit integer
+// between 512 and 1048576 indicating the size in bytes).
+// 4. The client sends RSS_SET_PATH, followed by a packet (as defined above)
+// containing a path (in the server's filesystem) or a wildcard pattern
+// indicating the files/directories the client wants to download.
+// The server then searches the given path. For every file found:
+// 5. The server sends RSS_CREATE_FILE, followed by a packet containing the
+// filename (filename only, without a path), followed by a series of
+// packets (called chunks) containing the file's contents. The size of
+// each chunk is the size set by the client in step 3, except for the
+// last chunk, which must be smaller.
+// For every directory found:
+// 6. The server sends RSS_CREATE_DIR, followed by a packet containing the
+// name of the directory to be created (directory name only, without a
+// path).
+// 7. The server sends files and directories located inside the directory
+// (using steps 5, 6, 8).
+// 8. The server sends RSS_LEAVE_DIR.
+// 9. The server sends RSS_DONE.
+// 10. Steps 4-9 are repeated as many times as necessary.
+// 11. The client disconnects.
+// If a critical error occurs, the server may send RSS_ERROR followed by a
+// packet containing an error message, and the connection is closed.
+// RSS_ERROR may be sent only when the client expects a msg.
+
+#define _WIN32_WINNT 0x0500
+
+#include <winsock2.h>
+#include <windows.h>
+#include <stdio.h>
+#include <stdarg.h>
+#include <shlwapi.h>
+
+#pragma comment(lib, "ws2_32.lib")
+#pragma comment(lib, "shlwapi.lib")
+
+#define TEXTBOX_LIMIT 262144
+
+// Constants for file transfer server
+#define RSS_MAGIC 0x525353
+#define RSS_OK 1
+#define RSS_ERROR 2
+#define RSS_UPLOAD 3
+#define RSS_DOWNLOAD 4
+#define RSS_SET_PATH 5
+#define RSS_CREATE_FILE 6
+#define RSS_CREATE_DIR 7
+#define RSS_LEAVE_DIR 8
+#define RSS_DONE 9
+
+// Globals
+int shell_port = 10022;
+int file_transfer_port = 10023;
+
+HWND hMainWindow = NULL;
+HWND hTextBox = NULL;
+HANDLE hTextBufferMutex = NULL;
+
+char text_buffer[8192] = {0};
+int text_size = 0;
+
+struct client_info {
+ SOCKET socket;
+ char addr_str[256];
+ int pid;
+ HWND hwnd;
+ HANDLE hJob;
+ HANDLE hChildOutputRead;
+ HANDLE hThreadChildToSocket;
+ char *chunk_buffer;
+ int chunk_size;
+};
+
+/*-----------------
+ * Shared functions
+ *-----------------*/
+
+void ExitOnError(const char *message, BOOL winsock = FALSE)
+{
+ LPVOID system_message;
+ char buffer[512];
+ int error_code;
+
+ if (winsock)
+ error_code = WSAGetLastError();
+ else
+ error_code = GetLastError();
+ WSACleanup();
+
+ FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER |
+ FORMAT_MESSAGE_FROM_SYSTEM,
+ NULL,
+ error_code,
+ MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
+ (LPTSTR)&system_message,
+ 0,
+ NULL);
+ sprintf(buffer,
+ "%s!\n"
+ "Error code = %d\n"
+ "Error message = %s",
+ message, error_code, (char *)system_message);
+ MessageBox(hMainWindow, buffer, "Error", MB_OK | MB_ICONERROR);
+
+ LocalFree(system_message);
+ ExitProcess(1);
+}
+
+void FlushTextBuffer()
+{
+ if (!text_size) return;
+ int len = GetWindowTextLength(hTextBox);
+ while (len > TEXTBOX_LIMIT - sizeof(text_buffer)) {
+ SendMessage(hTextBox, EM_SETSEL, 0, TEXTBOX_LIMIT * 1/4);
+ SendMessage(hTextBox, EM_REPLACESEL, FALSE, (LPARAM)"...");
+ len = GetWindowTextLength(hTextBox);
+ }
+ SendMessage(hTextBox, EM_SETSEL, len, len);
+ SendMessage(hTextBox, EM_REPLACESEL, FALSE, (LPARAM)text_buffer);
+ text_buffer[0] = 0;
+ text_size = 0;
+}
+
+void AppendMessage(const char *message, ...)
+{
+ va_list args;
+ char str[512] = {0};
+
+ va_start(args, message);
+ vsnprintf(str, sizeof(str) - 3, message, args);
+ va_end(args);
+ strcat(str, "\r\n");
+ int len = strlen(str);
+
+ WaitForSingleObject(hTextBufferMutex, INFINITE);
+ if (text_size + len + 1 > sizeof(text_buffer))
+ FlushTextBuffer();
+ strcpy(text_buffer + text_size, str);
+ text_size += len;
+ ReleaseMutex(hTextBufferMutex);
+}
+
+// Flush the text buffer every 250 ms
+DWORD WINAPI UpdateTextBox(LPVOID client_info_ptr)
+{
+ while (1) {
+ Sleep(250);
+ WaitForSingleObject(hTextBufferMutex, INFINITE);
+ FlushTextBuffer();
+ ReleaseMutex(hTextBufferMutex);
+ }
+ return 0;
+}
+
+void FormatStringForPrinting(char *dst, const char *src, int size)
+{
+ int j = 0;
+
+ for (int i = 0; i < size && src[i]; i++) {
+ if (src[i] == '\n') {
+ dst[j++] = '\\';
+ dst[j++] = 'n';
+ } else if (src[i] == '\r') {
+ dst[j++] = '\\';
+ dst[j++] = 'r';
+ } else if (src[i] == '\t') {
+ dst[j++] = '\\';
+ dst[j++] = 't';
+ } else if (src[i] == '\\') {
+ dst[j++] = '\\';
+ dst[j++] = '\\';
+ } else dst[j++] = src[i];
+ }
+ dst[j] = 0;
+}
+
+SOCKET PrepareListenSocket(int port)
+{
+ sockaddr_in addr;
+ linger l;
+ int result;
+
+ // Create socket
+ SOCKET ListenSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
+ if (ListenSocket == INVALID_SOCKET)
+ ExitOnError("Socket creation failed", TRUE);
+
+ // Enable lingering
+ l.l_linger = 10;
+ l.l_onoff = 1;
+ setsockopt(ListenSocket, SOL_SOCKET, SO_LINGER, (char *)&l, sizeof(l));
+
+ // Bind the socket
+ addr.sin_family = AF_INET;
+ addr.sin_addr.s_addr = htonl(INADDR_ANY);
+ addr.sin_port = htons(port);
+
+ result = bind(ListenSocket, (sockaddr *)&addr, sizeof(addr));
+ if (result == SOCKET_ERROR)
+ ExitOnError("bind failed", TRUE);
+
+ // Start listening for incoming connections
+ result = listen(ListenSocket, SOMAXCONN);
+ if (result == SOCKET_ERROR)
+ ExitOnError("listen failed", TRUE);
+
+ return ListenSocket;
+}
+
+client_info* Accept(SOCKET ListenSocket)
+{
+ sockaddr_in addr;
+ int addrlen = sizeof(addr);
+
+ // Accept the connection
+ SOCKET socket = accept(ListenSocket, (sockaddr *)&addr, &addrlen);
+ if (socket == INVALID_SOCKET) {
+ if (WSAGetLastError() == WSAEINTR)
+ return NULL;
+ else
+ ExitOnError("accept failed", TRUE);
+ }
+
+ // Allocate a new client_info struct
+ client_info *ci = (client_info *)calloc(1, sizeof(client_info));
+ if (!ci)
+ ExitOnError("Could not allocate client_info struct");
+ // Populate the new struct
+ ci->socket = socket;
+ const char *address = inet_ntoa(addr.sin_addr);
+ if (!address) address = "unknown";
+ sprintf(ci->addr_str, "%s:%d", address, addr.sin_port);
+
+ return ci;
+}
+
+// Read a given number of bytes into a buffer
+BOOL Receive(SOCKET socket, char *buffer, int len)
+{
+ while (len > 0) {
+ int bytes_received = recv(socket, buffer, len, 0);
+ if (bytes_received <= 0)
+ return FALSE;
+ buffer += bytes_received;
+ len -= bytes_received;
+ }
+ return TRUE;
+}
+
+// Send a given number of bytes from a buffer
+BOOL Send(SOCKET socket, const char *buffer, int len)
+{
+ while (len > 0) {
+ int bytes_sent = send(socket, buffer, len, 0);
+ if (bytes_sent <= 0)
+ return FALSE;
+ buffer += bytes_sent;
+ len -= bytes_sent;
+ }
+ return TRUE;
+}
+
+/*-------------
+ * Shell server
+ *-------------*/
+
+DWORD WINAPI ChildToSocket(LPVOID client_info_ptr)
+{
+ client_info *ci = (client_info *)client_info_ptr;
+ char buffer[1024];
+ DWORD bytes_read;
+
+ while (1) {
+ // Read data from the child's STDOUT/STDERR pipes
+ if (!ReadFile(ci->hChildOutputRead,
+ buffer, sizeof(buffer),
+ &bytes_read, NULL) || !bytes_read) {
+ if (GetLastError() == ERROR_BROKEN_PIPE)
+ break; // Pipe done -- normal exit path
+ else
+ ExitOnError("ReadFile failed"); // Something bad happened
+ }
+ // Send data to the client
+ Send(ci->socket, buffer, bytes_read);
+ }
+
+ AppendMessage("Child exited");
+ closesocket(ci->socket);
+ return 0;
+}
+
+DWORD WINAPI SocketToChild(LPVOID client_info_ptr)
+{
+ client_info *ci = (client_info *)client_info_ptr;
+ char buffer[256], formatted_buffer[768];
+ int bytes_received;
+
+ AppendMessage("Shell server: new client connected (%s)", ci->addr_str);
+
+ while (1) {
+ // Receive data from the socket
+ ZeroMemory(buffer, sizeof(buffer));
+ bytes_received = recv(ci->socket, buffer, sizeof(buffer), 0);
+ if (bytes_received <= 0)
+ break;
+ // Report the data received
+ FormatStringForPrinting(formatted_buffer, buffer, sizeof(buffer));
+ AppendMessage("Client (%s) entered text: \"%s\"",
+ ci->addr_str, formatted_buffer);
+ // Send the data as a series of WM_CHAR messages to the console window
+ for (int i = 0; i < bytes_received; i++) {
+ SendMessage(ci->hwnd, WM_CHAR, buffer[i], 0);
+ SendMessage(ci->hwnd, WM_SETFOCUS, 0, 0);
+ }
+ }
+
+ AppendMessage("Shell server: client disconnected (%s)", ci->addr_str);
+
+ // Attempt to terminate the child's process tree:
+ // Using taskkill (where available)
+ sprintf(buffer, "taskkill /PID %d /T /F", ci->pid);
+ system(buffer);
+ // .. and using TerminateJobObject()
+ TerminateJobObject(ci->hJob, 0);
+ // Wait for the ChildToSocket thread to terminate
+ WaitForSingleObject(ci->hThreadChildToSocket, 10000);
+ // In case the thread refuses to exit, terminate it
+ TerminateThread(ci->hThreadChildToSocket, 0);
+ // Close the socket
+ closesocket(ci->socket);
+
+ // Free resources
+ CloseHandle(ci->hJob);
+ CloseHandle(ci->hThreadChildToSocket);
+ CloseHandle(ci->hChildOutputRead);
+ free(ci);
+
+ AppendMessage("SocketToChild thread exited");
+ return 0;
+}
+
+void PrepAndLaunchRedirectedChild(client_info *ci,
+ HANDLE hChildStdOut,
+ HANDLE hChildStdErr)
+{
+ PROCESS_INFORMATION pi;
+ STARTUPINFO si;
+
+ // Allocate a new console for the child
+ HWND hwnd = GetForegroundWindow();
+ FreeConsole();
+ AllocConsole();
+ ShowWindow(GetConsoleWindow(), SW_HIDE);
+ if (hwnd)
+ SetForegroundWindow(hwnd);
+
+ // Set up the start up info struct.
+ ZeroMemory(&si, sizeof(STARTUPINFO));
+ si.cb = sizeof(STARTUPINFO);
+ si.dwFlags = STARTF_USESTDHANDLES | STARTF_USESHOWWINDOW;
+ si.hStdOutput = hChildStdOut;
+ si.hStdInput = GetStdHandle(STD_INPUT_HANDLE);
+ si.hStdError = hChildStdErr;
+ // Use this if you want to hide the child:
+ si.wShowWindow = SW_HIDE;
+ // Note that dwFlags must include STARTF_USESHOWWINDOW if you want to
+ // use the wShowWindow flags.
+
+ // Launch the process that you want to redirect.
+ if (!CreateProcess(NULL, "cmd.exe", NULL, NULL, TRUE,
+ 0, NULL, "C:\\", &si, &pi))
+ ExitOnError("CreateProcess failed");
+
+ // Close any unnecessary handles.
+ if (!CloseHandle(pi.hThread))
+ ExitOnError("CloseHandle failed");
+
+ // Keep the process ID
+ ci->pid = pi.dwProcessId;
+ // Assign the process to a newly created JobObject
+ ci->hJob = CreateJobObject(NULL, NULL);
+ AssignProcessToJobObject(ci->hJob, pi.hProcess);
+ // Keep the console window's handle
+ ci->hwnd = GetConsoleWindow();
+
+ // Detach from the child's console
+ FreeConsole();
+}
+
+void SpawnSession(client_info *ci)
+{
+ HANDLE hOutputReadTmp, hOutputRead, hOutputWrite;
+ HANDLE hErrorWrite;
+ SECURITY_ATTRIBUTES sa;
+
+ // Set up the security attributes struct.
+ sa.nLength = sizeof(SECURITY_ATTRIBUTES);
+ sa.lpSecurityDescriptor = NULL;
+ sa.bInheritHandle = TRUE;
+
+ // Create the child output pipe.
+ if (!CreatePipe(&hOutputReadTmp, &hOutputWrite, &sa, 0))
+ ExitOnError("CreatePipe failed");
+
+ // Create a duplicate of the output write handle for the std error
+ // write handle. This is necessary in case the child application
+ // closes one of its std output handles.
+ if (!DuplicateHandle(GetCurrentProcess(), hOutputWrite,
+ GetCurrentProcess(), &hErrorWrite, 0,
+ TRUE, DUPLICATE_SAME_ACCESS))
+ ExitOnError("DuplicateHandle failed");
+
+ // Create new output read handle and the input write handles. Set
+ // the Properties to FALSE. Otherwise, the child inherits the
+ // properties and, as a result, non-closeable handles to the pipes
+ // are created.
+ if (!DuplicateHandle(GetCurrentProcess(), hOutputReadTmp,
+ GetCurrentProcess(),
+ &hOutputRead, // Address of new handle.
+ 0, FALSE, // Make it uninheritable.
+ DUPLICATE_SAME_ACCESS))
+ ExitOnError("DuplicateHandle failed");
+
+ // Close inheritable copies of the handles you do not want to be
+ // inherited.
+ if (!CloseHandle(hOutputReadTmp))
+ ExitOnError("CloseHandle failed");
+
+ PrepAndLaunchRedirectedChild(ci, hOutputWrite, hErrorWrite);
+
+ ci->hChildOutputRead = hOutputRead;
+
+ // Close pipe handles (do not continue to modify the parent).
+ // You need to make sure that no handles to the write end of the
+ // output pipe are maintained in this process or else the pipe will
+ // not close when the child process exits and the ReadFile will hang.
+ if (!CloseHandle(hOutputWrite)) ExitOnError("CloseHandle failed");
+ if (!CloseHandle(hErrorWrite)) ExitOnError("CloseHandle failed");
+}
+
+DWORD WINAPI ShellListenThread(LPVOID param)
+{
+ HANDLE hThread;
+
+ SOCKET ListenSocket = PrepareListenSocket(shell_port);
+
+ // Inform the user
+ AppendMessage("Shell server: waiting for clients to connect...");
+
+ while (1) {
+ client_info *ci = Accept(ListenSocket);
+ if (!ci) break;
+ // Under heavy load, spawning cmd.exe might take a while, so tell the
+ // client to be patient
+ const char *message = "Please wait...\r\n";
+ Send(ci->socket, message, strlen(message));
+ // Spawn a new redirected cmd.exe process
+ SpawnSession(ci);
+ // Start transferring data from the child process to the client
+ hThread = CreateThread(NULL, 0, ChildToSocket, (LPVOID)ci, 0, NULL);
+ if (!hThread)
+ ExitOnError("Could not create ChildToSocket thread");
+ ci->hThreadChildToSocket = hThread;
+ // ... and from the client to the child process
+ hThread = CreateThread(NULL, 0, SocketToChild, (LPVOID)ci, 0, NULL);
+ if (!hThread)
+ ExitOnError("Could not create SocketToChild thread");
+ }
+
+ return 0;
+}
+
+/*---------------------
+ * File transfer server
+ *---------------------*/
+
+int ReceivePacket(SOCKET socket, char *buffer, DWORD max_size)
+{
+ DWORD packet_size = 0;
+
+ if (!Receive(socket, (char *)&packet_size, 4))
+ return -1;
+ if (packet_size > max_size)
+ return -1;
+ if (!Receive(socket, buffer, packet_size))
+ return -1;
+
+ return packet_size;
+}
+
+int ReceiveStrPacket(SOCKET socket, char *buffer, DWORD max_size)
+{
+ memset(buffer, 0, max_size);
+ return ReceivePacket(socket, buffer, max_size - 1);
+}
+
+BOOL SendPacket(SOCKET socket, const char *buffer, DWORD len)
+{
+ if (!Send(socket, (char *)&len, 4))
+ return FALSE;
+ return Send(socket, buffer, len);
+}
+
+BOOL SendMsg(SOCKET socket, DWORD msg)
+{
+ return Send(socket, (char *)&msg, 4);
+}
+
+// Send data from a file
+BOOL SendFileChunks(client_info *ci, const char *filename)
+{
+ FILE *fp = fopen(filename, "rb");
+ if (!fp) return FALSE;
+
+ while (1) {
+ int bytes_read = fread(ci->chunk_buffer, 1, ci->chunk_size, fp);
+ if (!SendPacket(ci->socket, ci->chunk_buffer, bytes_read))
+ break;
+ if (bytes_read < ci->chunk_size) {
+ if (ferror(fp))
+ break;
+ else {
+ fclose(fp);
+ return TRUE;
+ }
+ }
+ }
+
+ fclose(fp);
+ return FALSE;
+}
+
+// Receive data into a file
+BOOL ReceiveFileChunks(client_info *ci, const char *filename)
+{
+ FILE *fp = fopen(filename, "wb");
+ if (!fp) return FALSE;
+
+ while (1) {
+ int bytes_received = ReceivePacket(ci->socket, ci->chunk_buffer,
+ ci->chunk_size);
+ if (bytes_received < 0)
+ break;
+ if (bytes_received > 0)
+ if (fwrite(ci->chunk_buffer, bytes_received, 1, fp) < 1)
+ break;
+ if (bytes_received < ci->chunk_size) {
+ fclose(fp);
+ return TRUE;
+ }
+ }
+
+ fclose(fp);
+ return FALSE;
+}
+
+BOOL ExpandPath(char *path, int max_size)
+{
+ char temp[512];
+ int result;
+
+ PathRemoveBackslash(path);
+ result = ExpandEnvironmentStrings(path, temp, sizeof(temp));
+ if (result == 0 || result > sizeof(temp))
+ return FALSE;
+ strncpy(path, temp, max_size - 1);
+ return TRUE;
+}
+
+int TerminateTransfer(client_info *ci, const char *message)
+{
+ AppendMessage(message);
+ AppendMessage("File transfer server: client disconnected (%s)",
+ ci->addr_str);
+ closesocket(ci->socket);
+ free(ci->chunk_buffer);
+ free(ci);
+ return 0;
+}
+
+int TerminateWithError(client_info *ci, const char *message)
+{
+ SendMsg(ci->socket, RSS_ERROR);
+ SendPacket(ci->socket, message, strlen(message));
+ return TerminateTransfer(ci, message);
+}
+
+int ReceiveThread(client_info *ci)
+{
+ char path[512], filename[512];
+ DWORD msg;
+
+ AppendMessage("Client (%s) wants to upload files", ci->addr_str);
+
+ while (1) {
+ if (!Receive(ci->socket, (char *)&msg, 4))
+ return TerminateTransfer(ci, "Could not receive further msgs");
+
+ switch (msg) {
+ case RSS_SET_PATH:
+ if (ReceiveStrPacket(ci->socket, path, sizeof(path)) < 0)
+ return TerminateWithError(ci,
+ "RSS_SET_PATH: could not receive path, or path too long");
+ AppendMessage("Client (%s) set path to %s", ci->addr_str, path);
+ if (!ExpandPath(path, sizeof(path)))
+ return TerminateWithError(ci,
+ "RSS_SET_PATH: error expanding environment strings");
+ break;
+
+ case RSS_CREATE_FILE:
+ if (ReceiveStrPacket(ci->socket, filename, sizeof(filename)) < 0)
+ return TerminateWithError(ci,
+ "RSS_CREATE_FILE: could not receive filename");
+ if (PathIsDirectory(path))
+ PathAppend(path, filename);
+ AppendMessage("Client (%s) is uploading %s", ci->addr_str, path);
+ if (!ReceiveFileChunks(ci, path))
+ return TerminateWithError(ci,
+ "RSS_CREATE_FILE: error receiving or writing file "
+ "contents");
+ PathAppend(path, "..");
+ break;
+
+ case RSS_CREATE_DIR:
+ if (ReceiveStrPacket(ci->socket, filename, sizeof(filename)) < 0)
+ return TerminateWithError(ci,
+ "RSS_CREATE_DIR: could not receive dirname");
+ if (PathIsDirectory(path))
+ PathAppend(path, filename);
+ AppendMessage("Entering dir %s", path);
+ if (PathFileExists(path)) {
+ if (!PathIsDirectory(path))
+ return TerminateWithError(ci,
+ "RSS_CREATE_DIR: path exists and is not a directory");
+ } else {
+ if (!CreateDirectory(path, NULL))
+ return TerminateWithError(ci,
+ "RSS_CREATE_DIR: could not create directory");
+ }
+ break;
+
+ case RSS_LEAVE_DIR:
+ PathAppend(path, "..");
+ AppendMessage("Returning to dir %s", path);
+ break;
+
+ case RSS_DONE:
+ if (!SendMsg(ci->socket, RSS_OK))
+ return TerminateTransfer(ci,
+ "RSS_DONE: could not send OK msg");
+ break;
+
+ default:
+ return TerminateWithError(ci, "Received unexpected msg");
+ }
+ }
+}
+
+// Given a path or a pattern with wildcards, send files or directory trees to
+// the client
+int SendFiles(client_info *ci, const char *pattern)
+{
+ char path[512];
+ WIN32_FIND_DATA ffd;
+
+ HANDLE hFind = FindFirstFile(pattern, &ffd);
+ if (hFind == INVALID_HANDLE_VALUE) {
+ // If a weird error occurred (like failure to list directory contents
+ // due to insufficient permissions) print a warning and continue.
+ if (GetLastError() != ERROR_FILE_NOT_FOUND)
+ AppendMessage("WARNING: FindFirstFile failed on pattern %s",
+ pattern);
+ return 1;
+ }
+
+ strncpy(path, pattern, sizeof(path) - 1);
+ PathAppend(path, "..");
+
+ do {
+ if (ffd.dwFileAttributes & FILE_ATTRIBUTE_REPARSE_POINT)
+ continue;
+ if (ffd.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {
+ // Directory
+ if (!strcmp(ffd.cFileName, ".") || !strcmp(ffd.cFileName, ".."))
+ continue;
+ PathAppend(path, ffd.cFileName);
+ AppendMessage("Entering dir %s", path);
+ PathAppend(path, "*");
+ if (!SendMsg(ci->socket, RSS_CREATE_DIR)) {
+ FindClose(hFind);
+ return TerminateTransfer(ci,
+ "Could not send RSS_CREATE_DIR msg");
+ }
+ if (!SendPacket(ci->socket, ffd.cFileName,
+ strlen(ffd.cFileName))) {
+ FindClose(hFind);
+ return TerminateTransfer(ci, "Could not send dirname");
+ }
+ if (!SendFiles(ci, path)) {
+ FindClose(hFind);
+ return 0;
+ }
+ if (!SendMsg(ci->socket, RSS_LEAVE_DIR)) {
+ FindClose(hFind);
+ return TerminateTransfer(ci,
+ "Could not send RSS_LEAVE_DIR msg");
+ }
+ PathAppend(path, "..");
+ PathAppend(path, "..");
+ AppendMessage("Returning to dir %s", path);
+ } else {
+ // File
+ PathAppend(path, ffd.cFileName);
+ AppendMessage("Client (%s) is downloading %s", ci->addr_str, path);
+ // Make sure the file is readable
+ FILE *fp = fopen(path, "rb");
+ if (fp) fclose(fp);
+ else {
+ AppendMessage("WARNING: could not read file %s", path);
+ PathAppend(path, "..");
+ continue;
+ }
+ if (!SendMsg(ci->socket, RSS_CREATE_FILE)) {
+ FindClose(hFind);
+ return TerminateTransfer(ci,
+ "Could not send RSS_CREATE_FILE msg");
+ }
+ if (!SendPacket(ci->socket, ffd.cFileName,
+ strlen(ffd.cFileName))) {
+ FindClose(hFind);
+ return TerminateTransfer(ci, "Could not send filename");
+ }
+ if (!SendFileChunks(ci, path)) {
+ FindClose(hFind);
+ return TerminateTransfer(ci, "Could not send file contents");
+ }
+ PathAppend(path, "..");
+ }
+ } while (FindNextFile(hFind, &ffd));
+
+ if (GetLastError() == ERROR_NO_MORE_FILES) {
+ FindClose(hFind);
+ return 1;
+ } else {
+ FindClose(hFind);
+ return TerminateWithError(ci, "FindNextFile failed");
+ }
+}
+
+int SendThread(client_info *ci)
+{
+ char pattern[512];
+ DWORD msg;
+
+ AppendMessage("Client (%s) wants to download files", ci->addr_str);
+
+ while (1) {
+ if (!Receive(ci->socket, (char *)&msg, 4))
+ return TerminateTransfer(ci, "Could not receive further msgs");
+
+ switch (msg) {
+ case RSS_SET_PATH:
+ if (ReceiveStrPacket(ci->socket, pattern, sizeof(pattern)) < 0)
+ return TerminateWithError(ci,
+ "RSS_SET_PATH: could not receive path, or path too long");
+ AppendMessage("Client (%s) asked for %s", ci->addr_str, pattern);
+ if (!ExpandPath(pattern, sizeof(pattern)))
+ return TerminateWithError(ci,
+ "RSS_SET_PATH: error expanding environment strings");
+ if (!SendFiles(ci, pattern))
+ return 0;
+ if (!SendMsg(ci->socket, RSS_DONE))
+ return TerminateTransfer(ci,
+ "RSS_SET_PATH: could not send RSS_DONE msg");
+ break;
+
+ default:
+ return TerminateWithError(ci, "Received unexpected msg");
+ }
+ }
+}
+
+DWORD WINAPI TransferThreadEntry(LPVOID client_info_ptr)
+{
+ client_info *ci = (client_info *)client_info_ptr;
+ DWORD msg;
+
+ AppendMessage("File transfer server: new client connected (%s)",
+ ci->addr_str);
+
+ if (!SendMsg(ci->socket, RSS_MAGIC))
+ return TerminateTransfer(ci, "Could not send greeting message");
+ if (!Receive(ci->socket, (char *)&ci->chunk_size, 4))
+ return TerminateTransfer(ci, "Error receiving chunk size");
+ AppendMessage("Client (%s) set chunk size to %d", ci->addr_str,
+ ci->chunk_size);
+ if (ci->chunk_size > 1048576 || ci->chunk_size < 512)
+ return TerminateWithError(ci, "Client set invalid chunk size");
+ if (!(ci->chunk_buffer = (char *)malloc(ci->chunk_size)))
+ return TerminateWithError(ci, "Memory allocation error");
+ if (!Receive(ci->socket, (char *)&msg, 4))
+ return TerminateTransfer(ci, "Error receiving msg");
+
+ if (msg == RSS_UPLOAD)
+ return ReceiveThread(ci);
+ else if (msg == RSS_DOWNLOAD)
+ return SendThread(ci);
+ return TerminateWithError(ci, "Received unexpected msg");
+}
+
+DWORD WINAPI FileTransferListenThread(LPVOID param)
+{
+ SOCKET ListenSocket = PrepareListenSocket(file_transfer_port);
+
+ // Inform the user
+ AppendMessage("File transfer server: waiting for clients to connect...");
+
+ while (1) {
+ client_info *ci = Accept(ListenSocket);
+ if (!ci) break;
+ if (!CreateThread(NULL, 0, TransferThreadEntry, (LPVOID)ci, 0, NULL))
+ ExitOnError("Could not create file transfer thread");
+ }
+
+ return 0;
+}
+
+/*--------------------
+ * WndProc and WinMain
+ *--------------------*/
+
+LRESULT CALLBACK WndProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam)
+{
+ RECT rect;
+ WSADATA wsaData;
+
+ switch (msg) {
+ case WM_CREATE:
+ // Create text box
+ GetClientRect(hwnd, &rect);
+ hTextBox = CreateWindowEx(WS_EX_CLIENTEDGE,
+ "EDIT", "",
+ WS_CHILD | WS_VISIBLE | WS_VSCROLL |
+ ES_MULTILINE | ES_AUTOVSCROLL,
+ 20, 20,
+ rect.right - 40,
+ rect.bottom - 40,
+ hwnd,
+ NULL,
+ GetModuleHandle(NULL),
+ NULL);
+ if (!hTextBox)
+ ExitOnError("Could not create text box");
+ // Set font
+ SendMessage(hTextBox, WM_SETFONT,
+ (WPARAM)GetStockObject(DEFAULT_GUI_FONT),
+ MAKELPARAM(FALSE, 0));
+ // Set size limit
+ SendMessage(hTextBox, EM_LIMITTEXT, TEXTBOX_LIMIT, 0);
+ // Create mutex for text buffer access
+ hTextBufferMutex = CreateMutex(NULL, FALSE, NULL);
+ // Create text box update thread
+ if (!CreateThread(NULL, 0, UpdateTextBox, NULL, 0, NULL))
+ ExitOnError("Could not create text box update thread");
+ // Initialize Winsock
+ if (WSAStartup(MAKEWORD(2, 2), &wsaData))
+ ExitOnError("Winsock initialization failed");
+ // Start the listening threads
+ if (!CreateThread(NULL, 0, ShellListenThread, NULL, 0, NULL))
+ ExitOnError("Could not create shell server listen thread");
+ if (!CreateThread(NULL, 0, FileTransferListenThread, NULL, 0, NULL))
+ ExitOnError("Could not create file transfer server listen thread");
+ break;
+
+ case WM_SIZE:
+ MoveWindow(hTextBox, 20, 20,
+ LOWORD(lParam) - 40, HIWORD(lParam) - 40, TRUE);
+ break;
+
+ case WM_DESTROY:
+ if (WSACleanup())
+ ExitOnError("WSACleanup failed");
+ PostQuitMessage(0);
+ break;
+
+ default:
+ return DefWindowProc(hwnd, msg, wParam, lParam);
+ }
+
+ return 0;
+}
+
+int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance,
+ LPSTR lpCmdLine, int nShowCmd)
+{
+ WNDCLASSEX wc;
+ MSG msg;
+ char title[256];
+
+ if (strlen(lpCmdLine))
+ sscanf(lpCmdLine, "%d %d", &shell_port, &file_transfer_port);
+
+ sprintf(title, "Remote Shell Server (listening on ports %d, %d)",
+ shell_port, file_transfer_port);
+
+ // Create the window class
+ wc.cbSize = sizeof(WNDCLASSEX);
+ wc.style = CS_HREDRAW | CS_VREDRAW;
+ wc.lpfnWndProc = WndProc;
+ wc.cbClsExtra = 0;
+ wc.cbWndExtra = 0;
+ wc.hInstance = hInstance;
+ wc.hIcon = LoadIcon(NULL, IDI_APPLICATION);
+ wc.hIconSm = LoadIcon(NULL, IDI_APPLICATION);
+ wc.hbrBackground = (HBRUSH)(COLOR_BTNFACE + 1);
+ wc.lpszMenuName = NULL;
+ wc.lpszClassName = "RemoteShellServerWindowClass";
+ wc.hCursor = LoadCursor(NULL, IDC_ARROW);
+
+ if (!RegisterClassEx(&wc))
+ ExitOnError("Could not register window class");
+
+ // Create the main window
+ hMainWindow =
+ CreateWindow("RemoteShellServerWindowClass", title,
+ WS_OVERLAPPEDWINDOW,
+ 20, 20, 600, 400,
+ NULL, NULL, hInstance, NULL);
+ if (!hMainWindow)
+ ExitOnError("Could not create window");
+
+ //ShowWindow(hMainWindow, SW_SHOWMINNOACTIVE);
+ ShowWindow(hMainWindow, SW_SHOW);
+ UpdateWindow(hMainWindow);
+
+ // Main message loop
+ while (GetMessage(&msg, NULL, 0, 0)) {
+ TranslateMessage(&msg);
+ DispatchMessage(&msg);
+ }
+
+ ExitProcess(0);
+}
--
1.5.4.1
^ permalink raw reply related [flat|nested] 3+ messages in thread
* [KVM-AUTOTEST PATCH v4] [RFC] KVM test: add python client for rss file transfer services
2010-07-04 13:42 [KVM-AUTOTEST PATCH v4] [RFC] KVM test: rss.cpp: add file transfer support Michael Goldish
@ 2010-07-04 13:42 ` Michael Goldish
2010-07-07 20:34 ` Lucas Meneghel Rodrigues
0 siblings, 1 reply; 3+ messages in thread
From: Michael Goldish @ 2010-07-04 13:42 UTC (permalink / raw)
To: autotest, kvm; +Cc: Michael Goldish
See details in docstrings in rss_file_transfer.py.
See protocol details in deps/rss.cpp.
Changes from v3:
- Protocol change: instead of sending a file as one big packet, send it in
multiple chunks. See details in deps/rss.cpp.
Changes from v2:
- Raise FileTransferNotFoundError if no files/dirs are transferred (due to
a bad path or wildcard pattern)
- Make all connection related errors in the base class raise
FileTransferConnectError
Changes from v1:
- Use glob() instead of iglob() (Python 2.4 doesn't like iglob())
- Change a few comments
Signed-off-by: Michael Goldish <mgoldish@redhat.com>
---
| 431 +++++++++++++++++++++++++++++++++
1 files changed, 431 insertions(+), 0 deletions(-)
create mode 100755 client/tests/kvm/rss_file_transfer.py
--git a/client/tests/kvm/rss_file_transfer.py b/client/tests/kvm/rss_file_transfer.py
new file mode 100755
index 0000000..c584397
--- /dev/null
+++ b/client/tests/kvm/rss_file_transfer.py
@@ -0,0 +1,431 @@
+#!/usr/bin/python
+"""
+Client for file transfer services offered by RSS (Remote Shell Server).
+
+@author: Michael Goldish (mgoldish@redhat.com)
+@copyright: 2008-2010 Red Hat Inc.
+"""
+
+import socket, struct, time, sys, os, glob
+
+# Globals
+CHUNKSIZE = 65536
+
+# Protocol message constants
+RSS_MAGIC = 0x525353
+RSS_OK = 1
+RSS_ERROR = 2
+RSS_UPLOAD = 3
+RSS_DOWNLOAD = 4
+RSS_SET_PATH = 5
+RSS_CREATE_FILE = 6
+RSS_CREATE_DIR = 7
+RSS_LEAVE_DIR = 8
+RSS_DONE = 9
+
+# See rss.cpp for protocol details.
+
+
+class FileTransferError(Exception):
+ pass
+
+
+class FileTransferConnectError(FileTransferError):
+ pass
+
+
+class FileTransferTimeoutError(FileTransferError):
+ pass
+
+
+class FileTransferProtocolError(FileTransferError):
+ pass
+
+
+class FileTransferSendError(FileTransferError):
+ pass
+
+
+class FileTransferServerError(FileTransferError):
+ pass
+
+
+class FileTransferNotFoundError(FileTransferError):
+ pass
+
+
+class FileTransferClient(object):
+ """
+ Connect to a RSS (remote shell server) and transfer files.
+ """
+
+ def __init__(self, address, port, timeout=10):
+ """
+ Connect to a server.
+
+ @param address: The server's address
+ @param port: The server's port
+ @param timeout: Time duration to wait for connection to succeed
+ @raise FileTransferConnectError: Raised if the connection fails
+ @raise FileTransferProtocolError: Raised if an incorrect magic number
+ is received
+ """
+ self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+ self._socket.settimeout(timeout)
+ try:
+ self._socket.connect((address, port))
+ except socket.error:
+ raise FileTransferConnectError("Could not connect to server")
+ try:
+ if self._receive_msg(timeout) != RSS_MAGIC:
+ raise FileTransferConnectError("Received wrong magic number")
+ except FileTransferTimeoutError:
+ raise FileTransferConnectError("Timeout expired while waiting to "
+ "receive magic number")
+ self._send(struct.pack("=i", CHUNKSIZE))
+
+
+ def __del__(self):
+ self.close()
+
+
+ def close(self):
+ """
+ Close the connection.
+ """
+ self._socket.close()
+
+
+ def _send(self, str):
+ try:
+ self._socket.sendall(str)
+ except socket.error:
+ raise FileTransferSendError("Could not send data to server")
+
+
+ def _receive(self, size, timeout=10):
+ strs = []
+ end_time = time.time() + timeout
+ while size > 0:
+ try:
+ self._socket.settimeout(max(0.0001, end_time - time.time()))
+ data = self._socket.recv(size)
+ except socket.timeout:
+ raise FileTransferTimeoutError("Timeout expired while "
+ "receiving data from server")
+ except socket.error:
+ raise FileTransferProtocolError("Error receiving data from "
+ "server")
+ if not data:
+ raise FileTransferProtocolError("Connection closed "
+ "unexpectedly")
+ strs.append(data)
+ size -= len(data)
+ return "".join(strs)
+
+
+ def _send_packet(self, str):
+ self._send(struct.pack("=I", len(str)))
+ self._send(str)
+
+
+ def _receive_packet(self, timeout=10):
+ size = struct.unpack("=I", self._receive(4))[0]
+ return self._receive(size, timeout)
+
+
+ def _send_file_chunks(self, filename, timeout=30):
+ f = open(filename, "rb")
+ try:
+ end_time = time.time() + timeout
+ while time.time() < end_time:
+ data = f.read(CHUNKSIZE)
+ self._send_packet(data)
+ if len(data) < CHUNKSIZE:
+ break
+ else:
+ raise FileTransferTimeoutError("Timeout expired while sending "
+ "file %s" % filename)
+ finally:
+ f.close()
+
+
+ def _receive_file_chunks(self, filename, timeout=30):
+ f = open(filename, "wb")
+ try:
+ end_time = time.time() + timeout
+ while True:
+ try:
+ data = self._receive_packet(end_time - time.time())
+ except FileTransferTimeoutError:
+ raise FileTransferTimeoutError("Timeout expired while "
+ "receiving file %s" %
+ filename)
+ except FileTransferProtocolError:
+ raise FileTransferProtocolError("Error receiving file %s" %
+ filename)
+ f.write(data)
+ if len(data) < CHUNKSIZE:
+ break
+ finally:
+ f.close()
+
+
+ def _send_msg(self, msg, timeout=10):
+ self._send(struct.pack("=I", msg))
+
+
+ def _receive_msg(self, timeout=10):
+ s = self._receive(4, timeout)
+ return struct.unpack("=I", s)[0]
+
+
+ def _handle_transfer_error(self):
+ # Save original exception
+ e = sys.exc_info()
+ try:
+ # See if we can get an error message
+ msg = self._receive_msg()
+ except FileTransferError:
+ # No error message -- re-raise original exception
+ raise e[0], e[1], e[2]
+ if msg == RSS_ERROR:
+ errmsg = self._receive_packet()
+ raise FileTransferServerError("Server said: %s" % errmsg)
+ raise e[0], e[1], e[2]
+
+
+class FileUploadClient(FileTransferClient):
+ """
+ Connect to a RSS (remote shell server) and upload files or directory trees.
+ """
+
+ def __init__(self, address, port, timeout=10):
+ """
+ Connect to a server.
+
+ @param address: The server's address
+ @param port: The server's port
+ @param timeout: Time duration to wait for connection to succeed
+ @raise FileTransferConnectError: Raised if the connection fails
+ @raise FileTransferProtocolError: Raised if an incorrect magic number
+ is received
+ @raise FileTransferSendError: Raised if the RSS_UPLOAD message cannot
+ be sent to the server
+ """
+ super(FileUploadClient, self).__init__(address, port, timeout)
+ self._send_msg(RSS_UPLOAD)
+
+
+ def _upload_file(self, path, end_time):
+ if os.path.isfile(path):
+ self._send_msg(RSS_CREATE_FILE)
+ self._send_packet(os.path.basename(path))
+ self._send_file_chunks(path, max(0, end_time - time.time()))
+ elif os.path.isdir(path):
+ self._send_msg(RSS_CREATE_DIR)
+ self._send_packet(os.path.basename(path))
+ for filename in os.listdir(path):
+ self._upload_file(os.path.join(path, filename), end_time)
+ self._send_msg(RSS_LEAVE_DIR)
+
+
+ def upload(self, src_pattern, dst_path, timeout=600):
+ """
+ Send files or directory trees to the server.
+ The semantics of src_pattern and dst_path are similar to those of scp.
+ For example, the following are OK:
+ src_pattern='/tmp/foo.txt', dst_path='C:\\'
+ (uploads a single file)
+ src_pattern='/usr/', dst_path='C:\\Windows\\'
+ (uploads a directory tree recursively)
+ src_pattern='/usr/*', dst_path='C:\\Windows\\'
+ (uploads all files and directory trees under /usr/)
+ The following is not OK:
+ src_pattern='/tmp/foo.txt', dst_path='C:\\Windows\\*'
+ (wildcards are only allowed in src_pattern)
+
+ @param src_pattern: A path or wildcard pattern specifying the files or
+ directories to send to the server
+ @param dst_path: A path in the server's filesystem where the files will
+ be saved
+ @param timeout: Time duration in seconds to wait for the transfer to
+ complete
+ @raise FileTransferTimeoutError: Raised if timeout expires
+ @raise FileTransferServerError: Raised if something goes wrong and the
+ server sends an informative error message to the client
+ @note: Other exceptions can be raised.
+ """
+ end_time = time.time() + timeout
+ try:
+ try:
+ self._send_msg(RSS_SET_PATH)
+ self._send_packet(dst_path)
+ matches = glob.glob(src_pattern)
+ for filename in matches:
+ self._upload_file(os.path.abspath(filename), end_time)
+ self._send_msg(RSS_DONE)
+ except FileTransferTimeoutError:
+ raise
+ except FileTransferError:
+ self._handle_transfer_error()
+ else:
+ # If nothing was transferred, raise an exception
+ if not matches:
+ raise FileTransferNotFoundError("Pattern %s does not "
+ "match any files or "
+ "directories" %
+ src_pattern)
+ # Look for RSS_OK or RSS_ERROR
+ msg = self._receive_msg(max(0, end_time - time.time()))
+ if msg == RSS_OK:
+ return
+ elif msg == RSS_ERROR:
+ errmsg = self._receive_packet()
+ raise FileTransferServerError("Server said: %s" % errmsg)
+ else:
+ # Neither RSS_OK nor RSS_ERROR found
+ raise FileTransferProtocolError("Received unexpected msg")
+ except:
+ # In any case, if the transfer failed, close the connection
+ self.close()
+ raise
+
+
+class FileDownloadClient(FileTransferClient):
+ """
+ Connect to a RSS (remote shell server) and download files or directory trees.
+ """
+
+ def __init__(self, address, port, timeout=10):
+ """
+ Connect to a server.
+
+ @param address: The server's address
+ @param port: The server's port
+ @param timeout: Time duration to wait for connection to succeed
+ @raise FileTransferConnectError: Raised if the connection fails
+ @raise FileTransferProtocolError: Raised if an incorrect magic number
+ is received
+ @raise FileTransferSendError: Raised if the RSS_UPLOAD message cannot
+ be sent to the server
+ """
+ super(FileDownloadClient, self).__init__(address, port, timeout)
+ self._send_msg(RSS_DOWNLOAD)
+
+
+ def download(self, src_pattern, dst_path, timeout=600):
+ """
+ Receive files or directory trees from the server.
+ The semantics of src_pattern and dst_path are similar to those of scp.
+ For example, the following are OK:
+ src_pattern='C:\\foo.txt', dst_path='/tmp'
+ (downloads a single file)
+ src_pattern='C:\\Windows', dst_path='/tmp'
+ (downloads a directory tree recursively)
+ src_pattern='C:\\Windows\\*', dst_path='/tmp'
+ (downloads all files and directory trees under C:\\Windows)
+ The following is not OK:
+ src_pattern='C:\\Windows', dst_path='/tmp/*'
+ (wildcards are only allowed in src_pattern)
+
+ @param src_pattern: A path or wildcard pattern specifying the files or
+ directories, in the server's filesystem, that will be sent to
+ the client
+ @param dst_path: A path in the local filesystem where the files will
+ be saved
+ @param timeout: Time duration in seconds to wait for the transfer to
+ complete
+ @raise FileTransferTimeoutError: Raised if timeout expires
+ @raise FileTransferServerError: Raised if something goes wrong and the
+ server sends an informative error message to the client
+ @note: Other exceptions can be raised.
+ """
+ dst_path = os.path.abspath(dst_path)
+ end_time = time.time() + timeout
+ file_count = 0
+ dir_count = 0
+ try:
+ try:
+ self._send_msg(RSS_SET_PATH)
+ self._send_packet(src_pattern)
+ except FileTransferError:
+ self._handle_transfer_error()
+ while True:
+ msg = self._receive_msg()
+ if msg == RSS_CREATE_FILE:
+ # Receive filename and file contents
+ filename = self._receive_packet()
+ if os.path.isdir(dst_path):
+ dst_path = os.path.join(dst_path, filename)
+ self._receive_file_chunks(
+ dst_path, max(0, end_time - time.time()))
+ dst_path = os.path.dirname(dst_path)
+ file_count += 1
+ elif msg == RSS_CREATE_DIR:
+ # Receive dirname and create the directory
+ dirname = self._receive_packet()
+ if os.path.isdir(dst_path):
+ dst_path = os.path.join(dst_path, dirname)
+ if not os.path.isdir(dst_path):
+ os.mkdir(dst_path)
+ dir_count += 1
+ elif msg == RSS_LEAVE_DIR:
+ # Return to parent dir
+ dst_path = os.path.dirname(dst_path)
+ elif msg == RSS_DONE:
+ # Transfer complete
+ if not file_count and not dir_count:
+ raise FileTransferNotFoundError("Pattern %s does not "
+ "match any files or "
+ "directories that "
+ "could be downloaded" %
+ src_pattern)
+ break
+ elif msg == RSS_ERROR:
+ # Receive error message and abort
+ errmsg = self._receive_packet()
+ raise FileTransferServerError("Server said: %s" % errmsg)
+ else:
+ # Unexpected msg
+ raise FileTransferProtocolError("Received unexpected msg")
+ except:
+ # In any case, if the transfer failed, close the connection
+ self.close()
+ raise
+
+
+def main():
+ import optparse
+
+ usage = "usage: %prog [options] address port src_pattern dst_path"
+ parser = optparse.OptionParser(usage=usage)
+ parser.add_option("-d", "--download",
+ action="store_true", dest="download",
+ help="download files from server")
+ parser.add_option("-u", "--upload",
+ action="store_true", dest="upload",
+ help="upload files to server")
+ parser.add_option("-t", "--timeout",
+ type="int", dest="timeout", default=3600,
+ help="transfer timeout")
+ options, args = parser.parse_args()
+ if options.download == options.upload:
+ parser.error("you must specify either -d or -u")
+ if len(args) != 4:
+ parser.error("incorrect number of arguments")
+ address, port, src_pattern, dst_path = args
+ port = int(port)
+
+ if options.download:
+ client = FileDownloadClient(address, port)
+ client.download(src_pattern, dst_path, timeout=options.timeout)
+ client.close()
+ elif options.upload:
+ client = FileUploadClient(address, port)
+ client.upload(src_pattern, dst_path, timeout=options.timeout)
+ client.close()
+
+
+if __name__ == "__main__":
+ main()
--
1.5.4.1
^ permalink raw reply related [flat|nested] 3+ messages in thread
* Re: [KVM-AUTOTEST PATCH v4] [RFC] KVM test: add python client for rss file transfer services
2010-07-04 13:42 ` [KVM-AUTOTEST PATCH v4] [RFC] KVM test: add python client for rss file transfer services Michael Goldish
@ 2010-07-07 20:34 ` Lucas Meneghel Rodrigues
0 siblings, 0 replies; 3+ messages in thread
From: Lucas Meneghel Rodrigues @ 2010-07-07 20:34 UTC (permalink / raw)
To: Michael Goldish; +Cc: autotest, kvm
On Sun, 2010-07-04 at 16:42 +0300, Michael Goldish wrote:
> See details in docstrings in rss_file_transfer.py.
> See protocol details in deps/rss.cpp.
I've gone through both rss.cpp code and rss_file_transfer.py code, and
must say it's very good. The only comment I'd make about this module is
that the command line syntax for it when using it as a stand alone
program could work similarly to utilities like rsync and scp, due to the
familiar model of usage, and also, the module/program could be called
rsscp instead of rss_file_transfer. Of course this is a minor comment,
and can be changed at a later time.
I'll commit this patchset shortly, together with a patch making
guest_test able to use file transfer already, as well as publish a new
winutils.iso.
> Changes from v3:
> - Protocol change: instead of sending a file as one big packet, send it in
> multiple chunks. See details in deps/rss.cpp.
>
> Changes from v2:
> - Raise FileTransferNotFoundError if no files/dirs are transferred (due to
> a bad path or wildcard pattern)
> - Make all connection related errors in the base class raise
> FileTransferConnectError
>
> Changes from v1:
> - Use glob() instead of iglob() (Python 2.4 doesn't like iglob())
> - Change a few comments
>
> Signed-off-by: Michael Goldish <mgoldish@redhat.com>
> ---
> client/tests/kvm/rss_file_transfer.py | 431 +++++++++++++++++++++++++++++++++
> 1 files changed, 431 insertions(+), 0 deletions(-)
> create mode 100755 client/tests/kvm/rss_file_transfer.py
>
> diff --git a/client/tests/kvm/rss_file_transfer.py b/client/tests/kvm/rss_file_transfer.py
> new file mode 100755
> index 0000000..c584397
> --- /dev/null
> +++ b/client/tests/kvm/rss_file_transfer.py
> @@ -0,0 +1,431 @@
> +#!/usr/bin/python
> +"""
> +Client for file transfer services offered by RSS (Remote Shell Server).
> +
> +@author: Michael Goldish (mgoldish@redhat.com)
> +@copyright: 2008-2010 Red Hat Inc.
> +"""
> +
> +import socket, struct, time, sys, os, glob
> +
> +# Globals
> +CHUNKSIZE = 65536
> +
> +# Protocol message constants
> +RSS_MAGIC = 0x525353
> +RSS_OK = 1
> +RSS_ERROR = 2
> +RSS_UPLOAD = 3
> +RSS_DOWNLOAD = 4
> +RSS_SET_PATH = 5
> +RSS_CREATE_FILE = 6
> +RSS_CREATE_DIR = 7
> +RSS_LEAVE_DIR = 8
> +RSS_DONE = 9
> +
> +# See rss.cpp for protocol details.
> +
> +
> +class FileTransferError(Exception):
> + pass
> +
> +
> +class FileTransferConnectError(FileTransferError):
> + pass
> +
> +
> +class FileTransferTimeoutError(FileTransferError):
> + pass
> +
> +
> +class FileTransferProtocolError(FileTransferError):
> + pass
> +
> +
> +class FileTransferSendError(FileTransferError):
> + pass
> +
> +
> +class FileTransferServerError(FileTransferError):
> + pass
> +
> +
> +class FileTransferNotFoundError(FileTransferError):
> + pass
> +
> +
> +class FileTransferClient(object):
> + """
> + Connect to a RSS (remote shell server) and transfer files.
> + """
> +
> + def __init__(self, address, port, timeout=10):
> + """
> + Connect to a server.
> +
> + @param address: The server's address
> + @param port: The server's port
> + @param timeout: Time duration to wait for connection to succeed
> + @raise FileTransferConnectError: Raised if the connection fails
> + @raise FileTransferProtocolError: Raised if an incorrect magic number
> + is received
> + """
> + self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
> + self._socket.settimeout(timeout)
> + try:
> + self._socket.connect((address, port))
> + except socket.error:
> + raise FileTransferConnectError("Could not connect to server")
> + try:
> + if self._receive_msg(timeout) != RSS_MAGIC:
> + raise FileTransferConnectError("Received wrong magic number")
> + except FileTransferTimeoutError:
> + raise FileTransferConnectError("Timeout expired while waiting to "
> + "receive magic number")
> + self._send(struct.pack("=i", CHUNKSIZE))
> +
> +
> + def __del__(self):
> + self.close()
> +
> +
> + def close(self):
> + """
> + Close the connection.
> + """
> + self._socket.close()
> +
> +
> + def _send(self, str):
> + try:
> + self._socket.sendall(str)
> + except socket.error:
> + raise FileTransferSendError("Could not send data to server")
> +
> +
> + def _receive(self, size, timeout=10):
> + strs = []
> + end_time = time.time() + timeout
> + while size > 0:
> + try:
> + self._socket.settimeout(max(0.0001, end_time - time.time()))
> + data = self._socket.recv(size)
> + except socket.timeout:
> + raise FileTransferTimeoutError("Timeout expired while "
> + "receiving data from server")
> + except socket.error:
> + raise FileTransferProtocolError("Error receiving data from "
> + "server")
> + if not data:
> + raise FileTransferProtocolError("Connection closed "
> + "unexpectedly")
> + strs.append(data)
> + size -= len(data)
> + return "".join(strs)
> +
> +
> + def _send_packet(self, str):
> + self._send(struct.pack("=I", len(str)))
> + self._send(str)
> +
> +
> + def _receive_packet(self, timeout=10):
> + size = struct.unpack("=I", self._receive(4))[0]
> + return self._receive(size, timeout)
> +
> +
> + def _send_file_chunks(self, filename, timeout=30):
> + f = open(filename, "rb")
> + try:
> + end_time = time.time() + timeout
> + while time.time() < end_time:
> + data = f.read(CHUNKSIZE)
> + self._send_packet(data)
> + if len(data) < CHUNKSIZE:
> + break
> + else:
> + raise FileTransferTimeoutError("Timeout expired while sending "
> + "file %s" % filename)
> + finally:
> + f.close()
> +
> +
> + def _receive_file_chunks(self, filename, timeout=30):
> + f = open(filename, "wb")
> + try:
> + end_time = time.time() + timeout
> + while True:
> + try:
> + data = self._receive_packet(end_time - time.time())
> + except FileTransferTimeoutError:
> + raise FileTransferTimeoutError("Timeout expired while "
> + "receiving file %s" %
> + filename)
> + except FileTransferProtocolError:
> + raise FileTransferProtocolError("Error receiving file %s" %
> + filename)
> + f.write(data)
> + if len(data) < CHUNKSIZE:
> + break
> + finally:
> + f.close()
> +
> +
> + def _send_msg(self, msg, timeout=10):
> + self._send(struct.pack("=I", msg))
> +
> +
> + def _receive_msg(self, timeout=10):
> + s = self._receive(4, timeout)
> + return struct.unpack("=I", s)[0]
> +
> +
> + def _handle_transfer_error(self):
> + # Save original exception
> + e = sys.exc_info()
> + try:
> + # See if we can get an error message
> + msg = self._receive_msg()
> + except FileTransferError:
> + # No error message -- re-raise original exception
> + raise e[0], e[1], e[2]
> + if msg == RSS_ERROR:
> + errmsg = self._receive_packet()
> + raise FileTransferServerError("Server said: %s" % errmsg)
> + raise e[0], e[1], e[2]
> +
> +
> +class FileUploadClient(FileTransferClient):
> + """
> + Connect to a RSS (remote shell server) and upload files or directory trees.
> + """
> +
> + def __init__(self, address, port, timeout=10):
> + """
> + Connect to a server.
> +
> + @param address: The server's address
> + @param port: The server's port
> + @param timeout: Time duration to wait for connection to succeed
> + @raise FileTransferConnectError: Raised if the connection fails
> + @raise FileTransferProtocolError: Raised if an incorrect magic number
> + is received
> + @raise FileTransferSendError: Raised if the RSS_UPLOAD message cannot
> + be sent to the server
> + """
> + super(FileUploadClient, self).__init__(address, port, timeout)
> + self._send_msg(RSS_UPLOAD)
> +
> +
> + def _upload_file(self, path, end_time):
> + if os.path.isfile(path):
> + self._send_msg(RSS_CREATE_FILE)
> + self._send_packet(os.path.basename(path))
> + self._send_file_chunks(path, max(0, end_time - time.time()))
> + elif os.path.isdir(path):
> + self._send_msg(RSS_CREATE_DIR)
> + self._send_packet(os.path.basename(path))
> + for filename in os.listdir(path):
> + self._upload_file(os.path.join(path, filename), end_time)
> + self._send_msg(RSS_LEAVE_DIR)
> +
> +
> + def upload(self, src_pattern, dst_path, timeout=600):
> + """
> + Send files or directory trees to the server.
> + The semantics of src_pattern and dst_path are similar to those of scp.
> + For example, the following are OK:
> + src_pattern='/tmp/foo.txt', dst_path='C:\\'
> + (uploads a single file)
> + src_pattern='/usr/', dst_path='C:\\Windows\\'
> + (uploads a directory tree recursively)
> + src_pattern='/usr/*', dst_path='C:\\Windows\\'
> + (uploads all files and directory trees under /usr/)
> + The following is not OK:
> + src_pattern='/tmp/foo.txt', dst_path='C:\\Windows\\*'
> + (wildcards are only allowed in src_pattern)
> +
> + @param src_pattern: A path or wildcard pattern specifying the files or
> + directories to send to the server
> + @param dst_path: A path in the server's filesystem where the files will
> + be saved
> + @param timeout: Time duration in seconds to wait for the transfer to
> + complete
> + @raise FileTransferTimeoutError: Raised if timeout expires
> + @raise FileTransferServerError: Raised if something goes wrong and the
> + server sends an informative error message to the client
> + @note: Other exceptions can be raised.
> + """
> + end_time = time.time() + timeout
> + try:
> + try:
> + self._send_msg(RSS_SET_PATH)
> + self._send_packet(dst_path)
> + matches = glob.glob(src_pattern)
> + for filename in matches:
> + self._upload_file(os.path.abspath(filename), end_time)
> + self._send_msg(RSS_DONE)
> + except FileTransferTimeoutError:
> + raise
> + except FileTransferError:
> + self._handle_transfer_error()
> + else:
> + # If nothing was transferred, raise an exception
> + if not matches:
> + raise FileTransferNotFoundError("Pattern %s does not "
> + "match any files or "
> + "directories" %
> + src_pattern)
> + # Look for RSS_OK or RSS_ERROR
> + msg = self._receive_msg(max(0, end_time - time.time()))
> + if msg == RSS_OK:
> + return
> + elif msg == RSS_ERROR:
> + errmsg = self._receive_packet()
> + raise FileTransferServerError("Server said: %s" % errmsg)
> + else:
> + # Neither RSS_OK nor RSS_ERROR found
> + raise FileTransferProtocolError("Received unexpected msg")
> + except:
> + # In any case, if the transfer failed, close the connection
> + self.close()
> + raise
> +
> +
> +class FileDownloadClient(FileTransferClient):
> + """
> + Connect to a RSS (remote shell server) and download files or directory trees.
> + """
> +
> + def __init__(self, address, port, timeout=10):
> + """
> + Connect to a server.
> +
> + @param address: The server's address
> + @param port: The server's port
> + @param timeout: Time duration to wait for connection to succeed
> + @raise FileTransferConnectError: Raised if the connection fails
> + @raise FileTransferProtocolError: Raised if an incorrect magic number
> + is received
> + @raise FileTransferSendError: Raised if the RSS_UPLOAD message cannot
> + be sent to the server
> + """
> + super(FileDownloadClient, self).__init__(address, port, timeout)
> + self._send_msg(RSS_DOWNLOAD)
> +
> +
> + def download(self, src_pattern, dst_path, timeout=600):
> + """
> + Receive files or directory trees from the server.
> + The semantics of src_pattern and dst_path are similar to those of scp.
> + For example, the following are OK:
> + src_pattern='C:\\foo.txt', dst_path='/tmp'
> + (downloads a single file)
> + src_pattern='C:\\Windows', dst_path='/tmp'
> + (downloads a directory tree recursively)
> + src_pattern='C:\\Windows\\*', dst_path='/tmp'
> + (downloads all files and directory trees under C:\\Windows)
> + The following is not OK:
> + src_pattern='C:\\Windows', dst_path='/tmp/*'
> + (wildcards are only allowed in src_pattern)
> +
> + @param src_pattern: A path or wildcard pattern specifying the files or
> + directories, in the server's filesystem, that will be sent to
> + the client
> + @param dst_path: A path in the local filesystem where the files will
> + be saved
> + @param timeout: Time duration in seconds to wait for the transfer to
> + complete
> + @raise FileTransferTimeoutError: Raised if timeout expires
> + @raise FileTransferServerError: Raised if something goes wrong and the
> + server sends an informative error message to the client
> + @note: Other exceptions can be raised.
> + """
> + dst_path = os.path.abspath(dst_path)
> + end_time = time.time() + timeout
> + file_count = 0
> + dir_count = 0
> + try:
> + try:
> + self._send_msg(RSS_SET_PATH)
> + self._send_packet(src_pattern)
> + except FileTransferError:
> + self._handle_transfer_error()
> + while True:
> + msg = self._receive_msg()
> + if msg == RSS_CREATE_FILE:
> + # Receive filename and file contents
> + filename = self._receive_packet()
> + if os.path.isdir(dst_path):
> + dst_path = os.path.join(dst_path, filename)
> + self._receive_file_chunks(
> + dst_path, max(0, end_time - time.time()))
> + dst_path = os.path.dirname(dst_path)
> + file_count += 1
> + elif msg == RSS_CREATE_DIR:
> + # Receive dirname and create the directory
> + dirname = self._receive_packet()
> + if os.path.isdir(dst_path):
> + dst_path = os.path.join(dst_path, dirname)
> + if not os.path.isdir(dst_path):
> + os.mkdir(dst_path)
> + dir_count += 1
> + elif msg == RSS_LEAVE_DIR:
> + # Return to parent dir
> + dst_path = os.path.dirname(dst_path)
> + elif msg == RSS_DONE:
> + # Transfer complete
> + if not file_count and not dir_count:
> + raise FileTransferNotFoundError("Pattern %s does not "
> + "match any files or "
> + "directories that "
> + "could be downloaded" %
> + src_pattern)
> + break
> + elif msg == RSS_ERROR:
> + # Receive error message and abort
> + errmsg = self._receive_packet()
> + raise FileTransferServerError("Server said: %s" % errmsg)
> + else:
> + # Unexpected msg
> + raise FileTransferProtocolError("Received unexpected msg")
> + except:
> + # In any case, if the transfer failed, close the connection
> + self.close()
> + raise
> +
> +
> +def main():
> + import optparse
> +
> + usage = "usage: %prog [options] address port src_pattern dst_path"
> + parser = optparse.OptionParser(usage=usage)
> + parser.add_option("-d", "--download",
> + action="store_true", dest="download",
> + help="download files from server")
> + parser.add_option("-u", "--upload",
> + action="store_true", dest="upload",
> + help="upload files to server")
> + parser.add_option("-t", "--timeout",
> + type="int", dest="timeout", default=3600,
> + help="transfer timeout")
> + options, args = parser.parse_args()
> + if options.download == options.upload:
> + parser.error("you must specify either -d or -u")
> + if len(args) != 4:
> + parser.error("incorrect number of arguments")
> + address, port, src_pattern, dst_path = args
> + port = int(port)
> +
> + if options.download:
> + client = FileDownloadClient(address, port)
> + client.download(src_pattern, dst_path, timeout=options.timeout)
> + client.close()
> + elif options.upload:
> + client = FileUploadClient(address, port)
> + client.upload(src_pattern, dst_path, timeout=options.timeout)
> + client.close()
> +
> +
> +if __name__ == "__main__":
> + main()
^ permalink raw reply [flat|nested] 3+ messages in thread
end of thread, other threads:[~2010-07-07 20:34 UTC | newest]
Thread overview: 3+ messages (download: mbox.gz / follow: Atom feed)
-- links below jump to the message on this page --
2010-07-04 13:42 [KVM-AUTOTEST PATCH v4] [RFC] KVM test: rss.cpp: add file transfer support Michael Goldish
2010-07-04 13:42 ` [KVM-AUTOTEST PATCH v4] [RFC] KVM test: add python client for rss file transfer services Michael Goldish
2010-07-07 20:34 ` Lucas Meneghel Rodrigues
This is an external index of several public inboxes,
see mirroring instructions on how to clone and mirror
all data and code used by this external index.