KaptureControlGui/kcg_remote_client.py

import socket
import argparse as ap
import readline


def connect(host, port):
    soc = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    soc.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    soc.connect((host, port))
    return soc

def send(string, host, port, p=True):
    soc = connect(host, port)
    soc.send(string)
    ret = ""
    if 'get_image' in string:
        f = open('/tmp/kcgimage.png', 'w+')
    while True:
        try:
            if 'get_image' in string:
                res = soc.recv(4096)
                if 'close' in res:
                    f.write(res[:-5])
                    f.close()
                    return "Done"
                f.write(res)
                continue
            soc.settimeout(20)
            res = soc.recv(1024)

            if res == 'close':
                if p:
                    return
                else:
                    # print ret
                    return ret
            # if '.' == res[0]:
            #     if p:
            #         print res[1:]
            #     else:
            #         ret += res[1:]
            else:
                if p:
                    print res
                else:
                    ret += res
        except socket.timeout:
            if not p:
                print "Timed Out"
            else:
                while True:
                    input = raw_input("Server takes long to respond, continue waiting? [y]es/[n]o ")
                    if input == 'y':
                        break
                    elif input == 'n':
                        return
                    else:
                        print 'wrong input'
    if 'get_image' in string:
        f.close()
        return
    if not p:
        return ret


def main():
    parse = ap.ArgumentParser("Remote Client for KCG - KAPTURE Control Gui")
    parse.add_argument("--port", "-p", type=int, required=True, help="Port")
    parse.add_argument("--host", "-H", type=str, required=True, help="Host")
    parse.add_argument("--command", "-c", type=str, help="Directly send this command and do not start interactive session.")
    args = parse.parse_args()

    if args.command:
        send(args.command, args.host, args.port)
    else:
        while True:
            input = raw_input("KCG > ")
            if input == 'stop':
                break
            send(input, args.host, args.port)

if __name__ == '__main__':
    main()