ProcDaemon/client.py at master · rwhelan/ProcDaemon · GitHub

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
import os
import select
import json
import time

from includes import asynclient
from proc import Process

print_calls = False

def clgr(fn):
    def wrapper(*args):
        if print_calls: print fn
        fn(*args)
    return wrapper

class TCPClient(asynclient):
    def __init__(self, g, sock, addr):
        self.g = g

        self.sock = sock
        self.addr = addr
        self.sockfd = self.sock.fileno()

        self.reg_fd(self.sockfd)

        self.stdoutbuf = ''
        self.stderrbuf = ''
        self.splerrbuf = ''

        self.havecmd = False

    @clgr
    def _sock_recv(self, fd, event):
        self.cmd = self.sock.recv(4096)
        if self.cmd and not self.havecmd:
            print time.ctime()+' - '+self.cmd
            self.havecmd = True

            self.proc = Process(self.cmd)

            self.reg_fd(self.proc.stdout)
            self.reg_fd(self.proc.stderr)
            self.reg_fd(self.proc.splerr)

            self.g['pids'][self.proc.pid] = self

        if not self.cmd:
            # If we're here, the client disconnected

            # Did they leave a proc running?
            if hasattr(self, 'proc'):
                self.un_reg_fd(self.proc.stdout)
                self.un_reg_fd(self.proc.stderr)
                self.un_reg_fd(self.proc.splerr)
                os.close(self.proc.stdout)
                os.close(self.proc.stderr)
                os.close(self.proc.splerr)
                os.kill(self.proc.pid, 9)

            self._close_net_sock()


    @clgr
    def _sock_hup(self, fd, event):
        raise NotImplementedError('Client._sock_hup')


    @clgr
    def _proc_recv(self, fd, event):
        try:
            if fd == self.proc.stdout:
                # Someday I'll use mutible data types...
                self.stdoutbuf += os.read(fd, 4096)
            elif fd == self.proc.stderr:
                self.stderrbuf += os.read(fd, 4096)
            elif fd == self.proc.splerr:
                self.splerrbuf += os.read(fd, 4096)

        except OSError, E:
            if E.errno != 9: # Bad file descriptor (fd closed)
                raise


    @clgr
    def _proc_hup(self, fd, event):
        self._proc_recv(fd, event)
        self.un_reg_fd(fd)
        os.close(fd)


    @clgr
    def _close_net_sock(self):
        if not hasattr(self, 'sockclosed'):
            self.un_reg_fd(self.sockfd)
            self.sock.close()
            self.sockclosed = True


    @clgr
    def handle_event(self, fd, event):
        if event & select.EPOLLIN:
            if fd == self.sockfd:
                self._sock_recv(fd, event)
            else:
                self._proc_recv(fd, event)

        if event & select.EPOLLHUP:
            if fd == self.sockfd:
                self._sock_hup(fd, event)
            else:
                self._proc_hup(fd, event)


    @clgr
    def proc_finish(self, code, res):
        if not hasattr(self, 'sockclosed'):

            self._proc_recv(self.proc.splerr, None)
            self._proc_recv(self.proc.stderr, None)
            self._proc_recv(self.proc.stdout, None)

            proc_usage = {}
            for rsrc in [i for i in dir(res) if i.startswith('ru_')]:
                proc_usage[rsrc] = getattr(res, rsrc)

            response = { 'stdout'     : self.stdoutbuf,
                         'stderr'     : self.stderrbuf,
                         'strerror'   : self.splerrbuf,
                         'exitcode'   : code >> 8,
                         'killsig'    : code & 255,
                         'resource'   : proc_usage,
                       }

            self.sock.send(json.dumps(response, indent = 4))
            self._close_net_sock()

        del self.g['pids'][self.proc.pid]