forked from diskoverdata/diskover-community
-
Notifications
You must be signed in to change notification settings - Fork 0
/
diskover_socket_server.py
285 lines (243 loc) · 9.45 KB
/
diskover_socket_server.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""diskover - Elasticsearch file system crawler
diskover is a file system crawler that index's
your file metadata into Elasticsearch.
See README.md or https://github.com/shirosaidev/diskover
for more information.
Copyright (C) Chris Park 2017
diskover is released under the Apache 2.0 license. See
LICENSE for the full license text.
"""
import diskover
import socket
import subprocess
try:
import queue as Queue
except ImportError:
import Queue
import threading
import uuid
import json
import time
import sys
# dict to hold socket tasks
socket_tasks = {}
# list of socket client
clientlist = []
def socket_thread_handler(threadnum, q, cliargs, logger, verbose):
"""This is the socket thread handler function.
It runs the command msg sent from client.
"""
BUFF = 1024
while True:
try:
c = q.get()
clientsock, addr = c
logger.debug(clientsock)
logger.debug(addr)
data = clientsock.recv(BUFF)
data = data.decode('utf-8')
logger.debug('Data:')
logger.debug(data)
if not data:
# close connection to client
clientsock.close()
logger.info("[thread-%s]: %s closed connection"
% (threadnum, str(addr)))
q.task_done()
continue
# check if ping msg
elif data == 'ping':
logger.info("[thread-%s]: Got ping from %s"
% (threadnum, str(addr)))
# send pong reply
message = b'pong'
clientsock.send(message)
logger.debug(message)
else:
# strip away any headers sent by curl
data = data.split('\r\n')[-1]
logger.info("[thread-%s]: Got command from %s"
% (threadnum, str(addr)))
# load json and store in dict
command_dict = json.loads(data)
logger.debug(command_dict)
# run command from json data
run_command(threadnum, command_dict, clientsock, cliargs, logger, verbose)
# close connection to client
clientsock.close()
logger.info("[thread-%s]: %s closed connection"
% (threadnum, str(addr)))
q.task_done()
except (ValueError, TypeError) as e:
logger.warning("[thread-%s]: Invalid JSON from %s: (%s)"
% (threadnum, str(addr), e))
message = b'{"msg": "error", "error": ' + e + b'}\n'
clientsock.send(message)
logger.debug(message)
# close connection to client
clientsock.close()
logger.info("[thread-%s]: %s closed connection"
% (threadnum, str(addr)))
q.task_done()
pass
except socket.error as e:
logger.error("[thread-%s]: Socket error (%s)" % (threadnum, e))
# close connection to client
clientsock.close()
logger.info("[thread-%s]: %s closed connection"
% (threadnum, str(addr)))
q.task_done()
pass
def start_socket_server(cliargs, logger, verbose):
"""This is the start socket server function.
It opens a socket and waits for remote commands.
"""
global clientlist
# set thread/connection limit
max_connections = 5
# Queue for socket threads
q = Queue.Queue(maxsize=max_connections)
try:
# create TCP socket object
serversock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
serversock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
host = diskover.config['listener_host'] # default is localhost
port = diskover.config['listener_port'] # default is 9999
# bind to port
serversock.bind((host, port))
# start listener
serversock.listen(max_connections)
# set up the threads and start them
for i in range(max_connections):
# create thread
t = threading.Thread(target=socket_thread_handler,
args=(i, q, cliargs, logger, verbose,))
t.daemon = True
t.start()
while True:
logger.info("Waiting for connection, listening on %s port %s TCP"
% (str(host), str(port)))
# establish connection
clientsock, addr = serversock.accept()
logger.debug(clientsock)
logger.debug(addr)
logger.info("Got a connection from %s" % str(addr))
# add client to list
client = [clientsock, addr]
clientlist.append(client)
# add task to Queue
q.put(client)
except socket.error as e:
serversock.close()
logger.error("Error opening socket (%s)" % e)
sys.exit(1)
except KeyboardInterrupt:
print('\nCtrl-c keyboard interrupt received, shutting down...')
q.join()
serversock.close()
sys.exit(0)
def run_command(threadnum, command_dict, clientsock, cliargs, logger, verbose):
"""This is the run command function.
It runs commands from the listener socket
using values in command_dict.
"""
global socket_tasks
global clientlist
# try to get index name from command or use from diskover config file
try:
index = str(command_dict['index'])
except KeyError:
index = str(diskover.config['index'])
pass
# try to get worker batch size from command or use default
try:
batchsize = str(command_dict['batchsize'])
except KeyError:
batchsize = str(cliargs['batchsize'])
pass
# try to get adaptive batch option from command or use default
try:
adaptivebatch = str(command_dict['adaptivebatch'])
except KeyError:
adaptivebatch = str(cliargs['adaptivebatch'])
pass
try:
action = command_dict['action']
pythonpath = diskover.config['python_path']
diskoverpath = diskover.config['diskover_path']
# set up command for different action
if action == 'crawl':
path = command_dict['path']
cmd = [pythonpath, diskoverpath, '-b', batchsize,
'-i', index, '-d', path, '-q']
elif action == 'finddupes':
cmd = [pythonpath, diskoverpath, '-b', batchsize,
'-i', index, '--finddupes', '-q']
elif action == 'hotdirs':
index2 = str(command_dict['index2'])
cmd = [pythonpath, diskoverpath, '-b', batchsize,
'-i', index, '--hotdirs', index2, '-q']
elif action == 'reindex':
try:
recursive = command_dict['recursive']
except KeyError:
recursive = 'false'
pass
path = command_dict['path']
if recursive == 'true':
cmd = [pythonpath, diskoverpath, '-b', batchsize,
'-i', index, '-d', path, '-R', '-q']
else:
cmd = [pythonpath, diskoverpath, '-b', batchsize,
'-i', index, '-d', path, '-r', '-q']
elif action == 'kill':
taskid = command_dict['taskid']
logger.info("[thread-%s]: Kill task message received! (taskid:%s)",
threadnum, taskid)
# do something here to kill task (future)
message = b'{"msg": "taskkilled"}\n'
clientsock.send(message)
return
else:
logger.warning("Unknown action")
message = b'{"error": "unknown action"}\n'
clientsock.send(message)
return
# add adaptive batch
if (adaptivebatch == "True" or adaptivebatch == "true"):
cmd.append('-a')
# run command using subprocess
starttime = time.time()
taskid = str(uuid.uuid4()).encode('utf-8')
# start process
process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
# add process to socket_tasks dict
socket_tasks[taskid] = process
message = b'{"msg": "taskstart", "taskid": "' + taskid + b'"}\n'
clientsock.send(message)
logger.info("[thread-%s]: Running command (taskid:%s)",
threadnum, taskid.decode('utf-8'))
logger.info(cmd)
output, error = process.communicate()
# send exit msg to client
exitcode = str(process.returncode).encode('utf-8')
logger.debug('Command output:')
logger.debug(output.decode('utf-8'))
logger.debug('Command error:')
logger.debug(error.decode('utf-8'))
elapsedtime = str(diskover.get_time(time.time() - starttime)).encode('utf-8')
logger.info("Finished command (taskid:%s), exit code: %s, elapsed time: %s"
% (taskid.decode('utf-8'), exitcode.decode('utf-8'), elapsedtime.decode('utf-8')))
message = b'{"msg": "taskfinish", "taskid": "%s", "exitcode": %s, "elapsedtime": "%s"}\n' \
% (taskid, exitcode, elapsedtime)
clientsock.send(message)
except ValueError:
logger.warning("Value error")
message = b'{"error": "value error"}\n'
clientsock.send(message)
pass
except socket.error as e:
logger.error("[thread-%s]: Socket error (%s)" % (threadnum, e))
pass