Peach.Publishers.file

1 2 ''' 3 Some default file publishers. Output generated data to a file, etc. 4 5 @author: Michael Eddington 6 @version: $Id: Peach.Publishers.file-pysrc.html 1138 2008-08-16 19:39:03Z meddingt $ 7 ''' 8 9 # 10 # Copyright (c) 2005-2008 Michael Eddington 11 # Copyright (c) 2004-2005 IOActive Inc. 12 # 13 # Permission is hereby granted, free of charge, to any person obtaining a copy 14 # of this software and associated documentation files (the "Software"), to deal 15 # in the Software without restriction, including without limitation the rights 16 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 17 # copies of the Software, and to permit persons to whom the Software is 18 # furnished to do so, subject to the following conditions: 19 # 20 # The above copyright notice and this permission notice shall be included in 21 # all copies or substantial portions of the Software. 22 # 23 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 24 # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 25 # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 26 # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 27 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 28 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 29 # SOFTWARE. 30 # 31 32 # Authors: 33 # Michael Eddington (mike@phed.org) 34 35 # $Id: Peach.Publishers.file-pysrc.html 1138 2008-08-16 19:39:03Z meddingt $ 36 37 import os, sys, time 38 from Peach.publisher import Publisher 39

40 -class FileWriter(Publisher):

41 ''' 42 Publishes generated data to a file. No concept of receaving data 43 yet. 44 ''' 45 46 _filename = None 47 _fd = None 48 _state = 0 # 0 -stoped; 1 -started 49

50 - def __init__(self, filename):

51 ''' 52 @type filename: string 53 @param filename: Filename to write to 54 ''' 55 self.setFilename(filename)

56

57 - def getFilename(self):

58 ''' 59 Get current filename. 60 61 @rtype: string 62 @return: current filename 63 ''' 64 return self._filename

65 - def setFilename(self, filename):

66 ''' 67 Set new filename. 68 69 @type filename: string 70 @param filename: Filename to set 71 ''' 72 self._filename = filename

73

74 - def start(self):

75 pass

76

77 - def connect(self):

78 if self._state == 1: 79 raise Exception('File::start(): Already started!') 80 81 if self._fd != None: 82 self._fd.close() 83 84 self.mkdir() 85 86 self._fd = open(self._filename, "w+b") 87 self._state = 1

88

89 - def stop(self):

90 self.close()

91

92 - def mkdir(self):

93 # lets try and create the folder this file lives in 94 delim = "" 95 96 if self._filename.find("\\") != -1: 97 delim = '\\' 98 elif self._filename.find("/") != -1: 99 delim = '/' 100 else: 101 return 102 103 # strip filename 104 try: 105 path = self._filename[: self._filename.rfind(delim) ] 106 os.mkdir(path) 107 except: 108 pass

109

110 - def close(self):

111 if self._state == 0: 112 return 113 114 self._fd.close() 115 self._fd = None 116 self._state = 0

117

118 - def send(self, data):

119 self._fd.write(data)

120

121 - def receive(self, size = None):

122 if size != None: 123 return self._fd.read(size) 124 125 return self._fd.read()

126

127 -class FileReader(Publisher):

128 ''' 129 Publishes generated data to a file. No concept of receaving data 130 yet. 131 ''' 132 133 _filename = None 134 _fd = None 135 _state = 0 # 0 -stoped; 1 -started 136

137 - def __init__(self, filename):

138 ''' 139 @type filename: string 140 @param filename: Filename to write to 141 ''' 142 self.setFilename(filename)

143

144 - def getFilename(self):

145 ''' 146 Get current filename. 147 148 @rtype: string 149 @return: current filename 150 ''' 151 return self._filename

152 - def setFilename(self, filename):

153 ''' 154 Set new filename. 155 156 @type filename: string 157 @param filename: Filename to set 158 ''' 159 self._filename = filename

160

161 - def start(self):

162 pass

163

164 - def connect(self):

165 if self._state == 1: 166 return 167 168 if self._fd != None: 169 self._fd.close() 170 171 self._fd = open(self._filename, "r+b") 172 self._state = 1

173

174 - def stop(self):

175 self.close()

176

177 - def close(self):

178 try: 179 if self._state == 0: 180 return 181 182 self._fd.close() 183 self._fd = None 184 self._state = 0 185 except: 186 pass

187

188 - def send(self, data):

189 self._fd.write(data)

190

191 - def receive(self, size = None):

192 if size != None: 193 return self._fd.read(size) 194 195 return self._fd.read()

196 197

198 -class FilePerIteration(FileWriter):

199 ''' 200 This publisher differs from File in that each round 201 will generate a new filename. Very handy for generating 202 bogus content (media files, etc). 203 ''' 204

205 - def __init__(self, filename):

206 ''' 207 @type filename: string 208 @param filename: Filename to write to should have a %d in it 209 someplace :) 210 ''' 211 FileWriter.__init__(self, filename) 212 213 self._roundCount = 0 214 self._origFilename = filename 215 self.setFilename(filename % self._roundCount) 216 self._closed = True

217

218 - def connect(self):

219 FileWriter.connect(self) 220 self._closed = False

221

222 - def stop(self):

223 self.close()

224

225 - def close(self):

226 FileWriter.close(self) 227 if not self._closed: 228 self._roundCount += 1 229 self.setFilename(self._origFilename % self._roundCount) 230 self._closed = True

231

232 - def send(self, data):

233 FileWriter.send(self, data)

234 235

236 -class FileWriterLauncher(Publisher):

237 ''' 238 Writes a file to disk and then launches a program. 239 240 To use, first use this publisher like the FileWriter 241 stream publisher. Close, than call a program (or two). 242 ''' 243 244 _filename = None 245 _fd = None 246 _state = 0 # 0 -stoped; 1 -started 247

248 - def __init__(self, filename):

249 ''' 250 @type filename: string 251 @param filename: Filename to write to 252 ''' 253 self.setFilename(filename)

254

255 - def getFilename(self):

256 ''' 257 Get current filename. 258 259 @rtype: string 260 @return: current filename 261 ''' 262 return self._filename

263 - def setFilename(self, filename):

264 ''' 265 Set new filename. 266 267 @type filename: string 268 @param filename: Filename to set 269 ''' 270 self._filename = filename

271

272 - def start(self):

273 pass

274

275 - def connect(self):

276 if self._state == 1: 277 raise Exception('File::start(): Already started!') 278 279 if self._fd != None: 280 self._fd.close() 281 282 self.mkdir() 283 284 self._fd = open(self._filename, "w+b") 285 self._state = 1

286

287 - def stop(self):

288 self.close()

289

290 - def mkdir(self):

291 # lets try and create the folder this file lives in 292 delim = "" 293 294 if self._filename.find("\\") != -1: 295 delim = '\\' 296 elif self._filename.find("/") != -1: 297 delim = '/' 298 else: 299 return 300 301 # strip filename 302 try: 303 path = self._filename[: self._filename.rfind(delim) ] 304 os.mkdir(path) 305 except: 306 pass

307

308 - def close(self):

309 if self._state == 0: 310 return 311 312 self._fd.close() 313 self._fd = None 314 self._state = 0

315

316 - def send(self, data):

317 self._fd.write(data)

318

319 - def receive(self, size = None):

320 if size != None: 321 return self._fd.read(size) 322 323 return self._fd.read()

324

325 - def call(self, method, args):

326 ''' 327 Launch program to consume file 328 329 @type method: string 330 @param method: Command to execute 331 @type args: array of objects 332 @param args: Arguments to pass 333 ''' 334 335 realArgs = [method, "/c", method] 336 for a in args: 337 realArgs.append(a) 338 339 print "Spawning %s" % method 340 return os.spawnv(os.P_WAIT, os.path.join( os.getenv('SystemRoot'), 'system32','cmd.exe'), realArgs)

341 342 try: 343 import win32gui, win32con, win32process 344 import sys,time, os, signal 345

346 - class FileWriterLauncherGui(Publisher):

347 ''' 348 Writes a file to disk and then launches a program. After 349 some defined amount of time we will try and close the GUI 350 application by sending WM_CLOSE than kill it. 351 352 To use, first use this publisher like the FileWriter 353 stream publisher. Close, than call a program (or two). 354 ''' 355 356 _filename = None 357 _fd = None 358 _state = 0 # 0 -stoped; 1 -started 359

360 - def __init__(self, filename, windowname):

361 ''' 362 @type filename: string 363 @param filename: Filename to write to 364 @type windowname: string 365 @param windowname: Partial window name to locate and kill 366 ''' 367 self.setFilename(filename) 368 self._windowName = windowname

369

370 - def getFilename(self):

371 ''' 372 Get current filename. 373 374 @rtype: string 375 @return: current filename 376 ''' 377 return self._filename

378

379 - def setFilename(self, filename):

380 ''' 381 Set new filename. 382 383 @type filename: string 384 @param filename: Filename to set 385 ''' 386 self._filename = filename

387

388 - def start(self):

389 pass

390

391 - def connect(self):

392 if self._state == 1: 393 raise Exception('File::start(): Already started!') 394 395 if self._fd != None: 396 self._fd.close() 397 398 self.mkdir() 399 400 # First lets rename the old file if there is one 401 402 try: 403 try: 404 os.unlink(self._filename + ".old.old") 405 except: 406 pass 407 408 try: 409 os.rename(self._filename + ".old", self._filename + ".old.old") 410 except: 411 pass 412 413 os.rename(self._filename, self._filename + ".old") 414 415 except: 416 pass 417 418 # If we can't open the file it might 419 # still be open. Lets retry a few times. 420 for i in range(10): 421 try: 422 self._fd = open(self._filename, "w+b") 423 break 424 except: 425 if i == 9: 426 raise 427 428 time.sleep(1) 429 430 self._state = 1

431

432 - def stop(self):

433 self.close()

434

435 - def mkdir(self):

436 # lets try and create the folder this file lives in 437 delim = "" 438 439 if self._filename.find("\\") != -1: 440 delim = '\\' 441 elif self._filename.find("/") != -1: 442 delim = '/' 443 else: 444 return 445 446 # strip filename 447 try: 448 path = self._filename[: self._filename.rfind(delim) ] 449 os.mkdir(path) 450 except: 451 pass

452

453 - def close(self):

454 if self._state == 0: 455 return 456 457 self._fd.close() 458 self._fd = None 459 self._state = 0

460

461 - def send(self, data):

462 self._fd.write(data)

463

464 - def receive(self, size = None):

465 if size != None: 466 return self._fd.read(size) 467 468 return self._fd.read()

469

470 - def call(self, method, args):

471 ''' 472 Launch program to consume file 473 474 @type method: string 475 @param method: Command to execute 476 @type args: array of objects 477 @param args: Arguments to pass 478 ''' 479 480 realArgs = [method, "/c", method] 481 for a in args: 482 realArgs.append(a) 483 484 #print "Spawning %s" % method 485 os.spawnv(os.P_NOWAIT, os.path.join( os.getenv('SystemRoot'), 'system32','cmd.exe'), realArgs) 486 487 # Wait 5 seconds 488 time.sleep(1) 489 490 self.closeApp(self._windowName)

491

492 - def enumCallback(hwnd, windowName):

493 ''' 494 Will get called by win32gui.EnumWindows, once for each 495 top level application window. 496 ''' 497 498 try: 499 500 # Get window title 501 title = win32gui.GetWindowText(hwnd) 502 503 # Is this our guy? 504 if title.find(windowName) == -1: 505 return 506 507 (threadId, processId) = win32process.GetWindowThreadProcessId(hwnd) 508 509 # Send WM_CLOSE message 510 try: 511 win32gui.PostMessage(hwnd, win32con.WM_CLOSE, 0, 0) 512 win32gui.PostQuitMessage(hwnd) 513 except: 514 pass 515 516 # Give it upto 5 sec 517 for i in range(100): 518 if win32process.GetExitCodeProcess(processId) != win32con.STILL_ACTIVE: 519 # Process exited already 520 return 521 522 time.sleep(0.25) 523 524 try: 525 # Kill application 526 win32process.TerminateProcess(processId, 0) 527 except: 528 pass 529 except: 530 pass

531 532 enumCallback = staticmethod(enumCallback) 533

534 - def closeApp(self, title):

535 ''' 536 Close Application by window title 537 ''' 538 #print "CloseApp: %s" % title 539 win32gui.EnumWindows(FileWriterLauncherGui.enumCallback, title)

540 541 except: 542 pass 543 544 # end 545

Source Code for Module Peach.Publishers.file