Fixed session agent build error
[public/netxms.git] / src / agent / nxsagent / main.cpp
1 /*
2 ** NetXMS Session Agent
3 ** Copyright (C) 2003-2016 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->serialize();
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 else
130 {
131 TCHAR buffer[1024];
132 _tprintf(_T("WTSQuerySessionInformation(WTSConnectState) failed (%s)\n"), GetSystemErrorText(GetLastError(), buffer, 1024));
133 }
134
135 msg.setField(VID_SESSION_STATE, sessionState);
136
137 TCHAR *sessionName;
138 if (WTSQuerySessionInformation(WTS_CURRENT_SERVER_HANDLE, sid, WTSWinStationName, &sessionName, &size))
139 {
140 if (*sessionName != 0)
141 {
142 msg.setField(VID_NAME, sessionName);
143 }
144 else
145 {
146 TCHAR buffer[256];
147 _sntprintf(buffer, 256, _T("%s-%d"), (sessionState == USER_SESSION_DISCONNECTED) ? _T("Disconnected") : ((sessionState == USER_SESSION_IDLE) ? _T("Idle") : _T("Session")), sid);
148 msg.setField(VID_NAME, buffer);
149 }
150 WTSFreeMemory(sessionName);
151 }
152 else
153 {
154 TCHAR buffer[1024];
155 _tprintf(_T("WTSQuerySessionInformation(WTSWinStationName) failed (%s)\n"), GetSystemErrorText(GetLastError(), buffer, 1024));
156 msg.setField(VID_NAME, _T("Console")); // assume console if session name cannot be read
157 }
158
159 TCHAR userName[256];
160 size = 256;
161 if (GetUserName(userName, &size))
162 {
163 msg.setField(VID_USER_NAME, userName);
164 }
165
166 SendMsg(&msg);
167 }
168
169 /**
170 * Process request from master agent
171 */
172 static void ProcessRequest(NXCPMessage *request)
173 {
174 NXCPMessage msg;
175
176 msg.setCode(CMD_REQUEST_COMPLETED);
177 msg.setId(request->getId());
178
179 switch(request->getCode())
180 {
181 case CMD_KEEPALIVE:
182 msg.setField(VID_RCC, ERR_SUCCESS);
183 break;
184 case CMD_TAKE_SCREENSHOT:
185 TakeScreenshot(&msg);
186 break;
187 default:
188 msg.setField(VID_RCC, ERR_UNKNOWN_COMMAND);
189 break;
190 }
191
192 SendMsg(&msg);
193 }
194
195 /**
196 * Message processing loop
197 */
198 static void ProcessMessages()
199 {
200 NXCPEncryptionContext *dummyCtx = NULL;
201 NXCPInitBuffer(&s_msgBuffer);
202 UINT32 rawMsgSize = 65536;
203 NXCP_MESSAGE *rawMsg = (NXCP_MESSAGE *)malloc(rawMsgSize);
204 while(true)
205 {
206 int err = RecvNXCPMessageEx(s_socket, &rawMsg, &s_msgBuffer, &rawMsgSize, &dummyCtx, NULL, 900000, 4 * 1024 * 1024);
207 if (err <= 0)
208 break;
209
210 // Check if message is too large
211 if (err == 1)
212 continue;
213
214 // Check for decryption failure
215 if (err == 2)
216 continue;
217
218 // Check for timeout
219 if (err == 3)
220 {
221 _tprintf(_T("Socket read timeout\n"));
222 break;
223 }
224
225 // Check that actual received packet size is equal to encoded in packet
226 if ((int)ntohl(rawMsg->size) != err)
227 {
228 _tprintf(_T("Actual message size doesn't match wSize value (%d,%d)\n"), err, ntohl(rawMsg->size));
229 continue; // Bad packet, wait for next
230 }
231
232 UINT16 flags = ntohs(rawMsg->flags);
233 if (!(flags & MF_BINARY))
234 {
235 NXCPMessage *msg = NXCPMessage::deserialize(rawMsg);
236 if (msg != NULL)
237 {
238 TCHAR msgCodeName[256];
239 _tprintf(_T("Received message %s\n"), NXCPMessageCodeName(msg->getCode(), msgCodeName));
240 ProcessRequest(msg);
241 delete msg;
242 }
243 else
244 {
245 _tprintf(_T("Message deserialization error\n"));
246 }
247 }
248 }
249 free(rawMsg);
250 }
251
252 #ifdef _WIN32
253
254 /**
255 * Window proc for event handling window
256 */
257 static LRESULT CALLBACK EventHandlerWndProc(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
258 {
259 switch(uMsg)
260 {
261 case WM_WTSSESSION_CHANGE:
262 _tprintf(_T(">> session change: %d\n"), wParam);
263 if ((wParam == WTS_CONSOLE_CONNECT) || (wParam == WTS_CONSOLE_DISCONNECT) ||
264 (wParam == WTS_REMOTE_CONNECT) || (wParam == WTS_REMOTE_DISCONNECT))
265 {
266 Login();
267 }
268 break;
269 default:
270 return DefWindowProc(hWnd, uMsg, wParam, lParam);
271 }
272 return 0;
273 }
274
275 /**
276 * Event handling thread
277 */
278 static THREAD_RESULT THREAD_CALL EventHandler(void *arg)
279 {
280 HINSTANCE hInstance = (HINSTANCE)GetModuleHandle(NULL);
281
282 WNDCLASS wc;
283 memset(&wc, 0, sizeof(WNDCLASS));
284 wc.lpfnWndProc = EventHandlerWndProc;
285 wc.hInstance = hInstance;
286 wc.cbWndExtra = 0;
287 wc.lpszClassName = _T("NetXMS_SessionAgent_Wnd");
288 if (RegisterClass(&wc) == 0)
289 {
290 _tprintf(_T("Call to RegisterClass() failed\n"));
291 return THREAD_OK;
292 }
293
294 HWND hWnd = CreateWindow(_T("NetXMS_SessionAgent_Wnd"), _T("NetXMS Session Agent"), 0, 0, 0, 0, 0, HWND_MESSAGE, NULL, hInstance, NULL);
295 if (hWnd == NULL)
296 {
297 _tprintf(_T("Cannot create window: %s"), GetSystemErrorText(GetLastError(), NULL, 0));
298 return THREAD_OK;
299 }
300
301 WTSRegisterSessionNotification(hWnd, NOTIFY_FOR_THIS_SESSION);
302
303 MSG msg;
304 while(GetMessage(&msg, NULL, 0, 0) > 0)
305 {
306 TranslateMessage(&msg);
307 DispatchMessage(&msg);
308 }
309
310 return THREAD_OK;
311 }
312
313 /**
314 * Get our own console window handle (an alternative to Microsoft's GetConsoleWindow)
315 */
316 static HWND GetConsoleHWND()
317 {
318 DWORD cpid = GetCurrentProcessId();
319 HWND hWnd = NULL;
320 while(true)
321 {
322 hWnd = FindWindowEx(NULL, hWnd, _T("ConsoleWindowClass"), NULL);
323 if (hWnd == NULL)
324 break;
325
326 DWORD wpid;
327 GetWindowThreadProcessId(hWnd, &wpid);
328 if (cpid == wpid)
329 break;
330 }
331
332 return hWnd;
333 }
334
335 #endif
336
337 /**
338 * Entry point
339 */
340 int main(int argc, char *argv[])
341 {
342 InitNetXMSProcess(true);
343
344 bool hideConsole = false;
345
346 int ch;
347 while((ch = getopt(argc, argv, "c:Hv")) != -1)
348 {
349 switch(ch)
350 {
351 case 'c': // config
352 break;
353 case 'H': // hide console
354 hideConsole = true;
355 break;
356 case 'v': // version
357 _tprintf(_T("NetXMS Session Agent Version ") NETXMS_VERSION_STRING _T(" Build ") NETXMS_VERSION_BUILD_STRING _T("\n"));
358 exit(0);
359 break;
360 case '?':
361 return 3;
362 }
363 }
364
365 #ifdef _WIN32
366 WSADATA wsaData;
367 int wrc = WSAStartup(MAKEWORD(2, 2), &wsaData);
368 if (wrc != 0)
369 {
370 _tprintf(_T("WSAStartup() failed"));
371 return 1;
372 }
373
374 ThreadCreate(EventHandler, 0, NULL);
375
376 if (hideConsole)
377 {
378 HWND hWnd = GetConsoleHWND();
379 if (hWnd != NULL)
380 ShowWindow(hWnd, SW_HIDE);
381 }
382 #endif
383
384 while(true)
385 {
386 if (!ConnectToMasterAgent())
387 {
388 ThreadSleep(30);
389 continue;
390 }
391
392 _tprintf(_T("*** Connected to master agent ***\n"));
393 Login();
394 ProcessMessages();
395 }
396 return 0;
397 }