da61a0820c67ed6ae63ae509ad08f365bff9e018
[public/netxms.git] / src / agent / nxsagent / main.cpp
1 /*
2 ** NetXMS Session Agent
3 ** Copyright (C) 2003-2014 Victor Kirhenshtein
4 **
5 ** This program is free software; you can redistribute it and/or modify
6 ** it under the terms of the GNU General Public License as published by
7 ** the Free Software Foundation; either version 2 of the License, or
8 ** (at your option) any later version.
9 **
10 ** This program is distributed in the hope that it will be usefu,,
11 ** but ITHOUT ANY WARRANTY; without even the implied warranty of
12 ** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13 ** GNU General Public License for more details.
14 **
15 ** You should have received a copy of the GNU General Public License
16 ** along with this program; if not, write to the Free Software
17 ** Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
18 **
19 ** File: main.cpp
20 **
21 **/
22
23 #include "nxsagent.h"
24
25 #if HAVE_GETOPT_H
26 #include <getopt.h>
27 #endif
28
29 /**
30 * Connection port
31 */
32 static UINT16 s_port = 28180;
33
34 /**
35 * Socket
36 */
37 static SOCKET s_socket = INVALID_SOCKET;
38
39 /**
40 * Socket lock
41 */
42 static MUTEX s_socketLock = MutexCreate();
43
44 /**
45 * Protocol buffer
46 */
47 static NXCP_BUFFER s_msgBuffer;
48
49 /**
50 * Connect to master agent
51 */
52 static bool ConnectToMasterAgent()
53 {
54 _tprintf(_T("Connecting to master agent...\n"));
55 s_socket = socket(AF_INET, SOCK_STREAM, 0);
56 if (s_socket == INVALID_SOCKET)
57 {
58 _tprintf(_T("Call to socket() failed\n"));
59 return false;
60 }
61
62 // Fill in address structure
63 struct sockaddr_in sa;
64 memset(&sa, 0, sizeof(sa));
65 sa.sin_family = AF_INET;
66 sa.sin_addr.s_addr = inet_addr("127.0.0.1");
67 sa.sin_port = htons(s_port);
68
69 // Connect to server
70 if (ConnectEx(s_socket, (struct sockaddr *)&sa, sizeof(sa), 5000) == -1)
71 {
72 _tprintf(_T("Cannot establish connection with master agent\n"));
73 closesocket(s_socket);
74 s_socket = INVALID_SOCKET;
75 return false;
76 }
77
78 return true;
79 }
80
81 /**
82 * Send message to master agent
83 */
84 static bool SendMsg(NXCPMessage *msg)
85 {
86 if (s_socket == INVALID_SOCKET)
87 return false;
88
89 NXCP_MESSAGE *rawMsg = msg->createMessage();
90 bool success = (SendEx(s_socket, rawMsg, ntohl(rawMsg->size), 0, s_socketLock) == ntohl(rawMsg->size));
91 free(rawMsg);
92 return success;
93 }
94
95 /**
96 * Send login message
97 */
98 static void Login()
99 {
100 NXCPMessage msg;
101 msg.setCode(CMD_LOGIN);
102
103 DWORD sid;
104 ProcessIdToSessionId(GetCurrentProcessId(), &sid);
105 msg.setField(VID_SESSION_ID, (UINT32)sid);
106
107 DWORD size;
108 WTS_CONNECTSTATE_CLASS *state;
109 INT16 sessionState = USER_SESSION_OTHER;
110 if (WTSQuerySessionInformation(WTS_CURRENT_SERVER_HANDLE, sid, WTSConnectState, (LPTSTR *)&state, &size))
111 {
112 switch(*state)
113 {
114 case WTSActive:
115 sessionState = USER_SESSION_ACTIVE;
116 break;
117 case WTSConnected:
118 sessionState = USER_SESSION_CONNECTED;
119 break;
120 case WTSDisconnected:
121 sessionState = USER_SESSION_DISCONNECTED;
122 break;
123 case WTSIdle:
124 sessionState = USER_SESSION_IDLE;
125 break;
126 }
127 WTSFreeMemory(state);
128 }
129
130 msg.setField(VID_SESSION_STATE, sessionState);
131
132 TCHAR *sessionName;
133 if (WTSQuerySessionInformation(WTS_CURRENT_SERVER_HANDLE, sid, WTSWinStationName, &sessionName, &size))
134 {
135 if (*sessionName != 0)
136 {
137 msg.setField(VID_NAME, sessionName);
138 }
139 else
140 {
141 TCHAR buffer[256];
142 _sntprintf(buffer, 256, _T("%s-%d"), (sessionState == USER_SESSION_DISCONNECTED) ? _T("Disconnected") : ((sessionState == USER_SESSION_IDLE) ? _T("Idle") : _T("Session")), sid);
143 msg.setField(VID_NAME, buffer);
144 }
145 WTSFreeMemory(sessionName);
146 }
147
148 TCHAR userName[256];
149 size = 256;
150 if (GetUserName(userName, &size))
151 {
152 msg.setField(VID_USER_NAME, userName);
153 }
154
155 SendMsg(&msg);
156 }
157
158 /**
159 * Process request from master agent
160 */
161 static void ProcessRequest(NXCPMessage *request)
162 {
163 NXCPMessage msg;
164
165 msg.setCode(CMD_REQUEST_COMPLETED);
166 msg.setId(request->getId());
167
168 switch(request->getCode())
169 {
170 case CMD_KEEPALIVE:
171 msg.setField(VID_RCC, ERR_SUCCESS);
172 break;
173 case CMD_TAKE_SCREENSHOT:
174 TakeScreenshot(&msg);
175 break;
176 default:
177 msg.setField(VID_RCC, ERR_UNKNOWN_COMMAND);
178 break;
179 }
180
181 SendMsg(&msg);
182 }
183
184 /**
185 * Message processing loop
186 */
187 static void ProcessMessages()
188 {
189 NXCPEncryptionContext *dummyCtx = NULL;
190 RecvNXCPMessage(0, NULL, &s_msgBuffer, 0, NULL, NULL, 0);
191 UINT32 rawMsgSize = 65536;
192 NXCP_MESSAGE *rawMsg = (NXCP_MESSAGE *)malloc(rawMsgSize);
193 while(1)
194 {
195 int err = RecvNXCPMessageEx(s_socket, &rawMsg, &s_msgBuffer, &rawMsgSize, &dummyCtx, NULL, 900000, 4 * 1024 * 1024);
196 if (err <= 0)
197 break;
198
199 // Check if message is too large
200 if (err == 1)
201 continue;
202
203 // Check for decryption failure
204 if (err == 2)
205 continue;
206
207 // Check for timeout
208 if (err == 3)
209 {
210 _tprintf(_T("Socket read timeout"));
211 break;
212 }
213
214 // Check that actual received packet size is equal to encoded in packet
215 if ((int)ntohl(rawMsg->size) != err)
216 {
217 _tprintf(_T("Actual message size doesn't match wSize value (%d,%d)"), err, ntohl(rawMsg->size));
218 continue; // Bad packet, wait for next
219 }
220
221 UINT16 flags = ntohs(rawMsg->flags);
222 if (!(flags & MF_BINARY))
223 {
224 NXCPMessage *msg = new NXCPMessage(rawMsg);
225 TCHAR msgCodeName[256];
226 _tprintf(_T("Received message %s\n"), NXCPMessageCodeName(msg->getCode(), msgCodeName));
227 ProcessRequest(msg);
228 delete msg;
229 }
230 }
231 free(rawMsg);
232 }
233
234 #ifdef _WIN32
235
236 /**
237 * Window proc for event handling window
238 */
239 static LRESULT CALLBACK EventHandlerWndProc(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
240 {
241 switch(uMsg)
242 {
243 case WM_WTSSESSION_CHANGE:
244 _tprintf(_T(">> session change: %d\n"), wParam);
245 if ((wParam == WTS_CONSOLE_CONNECT) || (wParam == WTS_CONSOLE_DISCONNECT) ||
246 (wParam == WTS_REMOTE_CONNECT) || (wParam == WTS_REMOTE_DISCONNECT))
247 {
248 Login();
249 }
250 break;
251 default:
252 return DefWindowProc(hWnd, uMsg, wParam, lParam);
253 }
254 return 0;
255 }
256
257 /**
258 * Event handling thread
259 */
260 static THREAD_RESULT THREAD_CALL EventHandler(void *arg)
261 {
262 HINSTANCE hInstance = (HINSTANCE)GetModuleHandle(NULL);
263
264 WNDCLASS wc;
265 memset(&wc, 0, sizeof(WNDCLASS));
266 wc.lpfnWndProc = EventHandlerWndProc;
267 wc.hInstance = hInstance;
268 wc.cbWndExtra = 0;
269 wc.lpszClassName = _T("NetXMS_SessionAgent_Wnd");
270 if (RegisterClass(&wc) == 0)
271 {
272 _tprintf(_T("Call to RegisterClass() failed\n"));
273 return THREAD_OK;
274 }
275
276 HWND hWnd = CreateWindow(_T("NetXMS_SessionAgent_Wnd"), _T("NetXMS Session Agent"), 0, 0, 0, 0, 0, HWND_MESSAGE, NULL, hInstance, NULL);
277 if (hWnd == NULL)
278 {
279 _tprintf(_T("Cannot create window: %s"), GetSystemErrorText(GetLastError(), NULL, 0));
280 return THREAD_OK;
281 }
282
283 WTSRegisterSessionNotification(hWnd, NOTIFY_FOR_THIS_SESSION);
284
285 MSG msg;
286 while(GetMessage(&msg, NULL, 0, 0) > 0)
287 {
288 TranslateMessage(&msg);
289 DispatchMessage(&msg);
290 }
291
292 return THREAD_OK;
293 }
294
295 /**
296 * Get our own console window handle (an alternative to Microsoft's GetConsoleWindow)
297 */
298 static HWND GetConsoleHWND()
299 {
300 HWND hWnd;
301 DWORD wpid, cpid;
302
303 cpid = GetCurrentProcessId();
304 while(1)
305 {
306 hWnd = FindWindowEx(NULL, NULL, _T("ConsoleWindowClass"), NULL);
307 if (hWnd == NULL)
308 break;
309
310 GetWindowThreadProcessId(hWnd, &wpid);
311 if (cpid == wpid)
312 break;
313 }
314
315 return hWnd;
316 }
317
318 #endif
319
320 /**
321 * Entry point
322 */
323 int main(int argc, char *argv[])
324 {
325 InitNetXMSProcess();
326
327 bool hideConsole = false;
328
329 int ch;
330 while((ch = getopt(argc, argv, "c:Hv")) != -1)
331 {
332 switch(ch)
333 {
334 case 'c': // config
335 break;
336 case 'H': // hide console
337 hideConsole = true;
338 break;
339 case 'v': // version
340 _tprintf(_T("NetXMS Session Agent Version ") NETXMS_VERSION_STRING _T(" Build ") NETXMS_VERSION_BUILD_STRING _T("\n"));
341 exit(0);
342 break;
343 case '?':
344 return 3;
345 }
346 }
347
348 #ifdef _WIN32
349 WSADATA wsaData;
350 int wrc = WSAStartup(MAKEWORD(2, 2), &wsaData);
351 if (wrc != 0)
352 {
353 _tprintf(_T("WSAStartup() failed"));
354 return 1;
355 }
356
357 ThreadCreate(EventHandler, 0, NULL);
358
359 if (hideConsole)
360 {
361 HWND hWnd = GetConsoleHWND();
362 if (hWnd != NULL)
363 ShowWindow(hWnd, SW_HIDE);
364 }
365 #endif
366
367 while(true)
368 {
369 if (!ConnectToMasterAgent())
370 {
371 ThreadSleep(30);
372 continue;
373 }
374
375 _tprintf(_T("*** Connected to master agent ***\n"));
376 Login();
377 ProcessMessages();
378 }
379 return 0;
380 }