-
Notifications
You must be signed in to change notification settings - Fork 0
/
willikins.py
executable file
·142 lines (115 loc) · 4.15 KB
/
willikins.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
#!/usr/bin/env python3
#
# willikins.py --- taking care of jobs and picking up the slack on Slack
#
# Willikins is a simple harness for posting information about running
# jobs to a Slack server. You supply Willikins with your username and
# the command that is to be run. Willikins executes it and sends some
# information about the command back to you.
import argparse
import json
import os
import shlex
import string
import subprocess
import time
from cgi import escape
from slackclient import SlackClient
from subprocess import Popen, PIPE
def format_output(rc, time):
"""
Formats the output that is used in sending the message to the client,
depending on the return code of the command.
"""
template = string.Template("""
${greeting}!\n
I have finished executing your job. Its return code is `$rc`, which
$comment.\n
Your command took `$time` to run. I am attaching its output to this
message. If the command created any files, you will have to inspect
them manually.
""")
return template.substitute( {
"greeting" : "Hello",
"rc" : rc,
"comment" : "looks good to me" if rc == 0 else "may indicate an error",
"time" : time
} )
def format_attachments(stdout, stderr):
"""
Formats the attachments of both `stdout` and `stderr` for a new
message to the client.
"""
attachments = []
for text, description in zip( [stdout, stderr], ["`stdout`", "`stderr`"]):
if not text:
continue
properties = {
"title" : description,
"fallback": description,
"text" : escape(text),
"mrkdwn" : "false",
}
attachments.append(properties)
return json.dumps(attachments)
if __name__ == "__main__":
######################################################################
# Parse command-line arguments
######################################################################
parser = argparse.ArgumentParser(description="Executes job with prejudice")
parser.add_argument("-a", "--attach", action="store_true", help="Attach command output instead of using a snippet")
parser.add_argument("-u", "--user", action="store" , type=str, help="ID of user to report to")
parser.add_argument("command" , metavar="COMMAND", action="store" , type=str, help="Command to run" )
arguments = parser.parse_args()
######################################################################
# Setting up the client
######################################################################
#
# Look for a direct channel to communicate with the user. If no such
# channel is found, the program exits.
token = os.getenv("SLACK_BOT_TOKEN")
user_id = arguments.user
command = arguments.command
attach = arguments.attach
sc = SlackClient(token)
im_channels = sc.api_call("im.list")
channel_id = None
for channel in im_channels['ims']:
if channel['user'] == user_id:
channel_id = channel['id']
assert channel_id
######################################################################
# Run the command
######################################################################
time_start = time.perf_counter()
commands = shlex.split(command)
p = Popen(shlex.split(command), stdout=PIPE, stderr=PIPE)
stdout, stderr = p.communicate()
rc = p.returncode
duration = "{:2f} s".format(time.perf_counter() - time_start)
message = format_output(rc, duration)
# TODO: make encoding configurable
stdout = stdout.decode("utf-8")
stderr = stderr.decode("utf-8")
if attach:
attachments = format_attachments(stdout, stderr)
else:
attachments = None
sc.api_call(
"chat.postMessage",
channel = channel_id,
text = message,
attachments = attachments
)
# User wants us to *upload* the output of the command instead of
# attaching it directly to the message.
if not attach:
for data, description in zip([stdout, stderr], ["stdout", "stderr"]):
if not data:
continue
sc.api_call(
"files.upload",
channels = channel_id,
content = data,
title = description
)