/
__init__.py
149 lines (124 loc) 路 6.28 KB
/
__init__.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
143
144
145
146
147
148
149
import logging
import random
import aiobotocore
import aiohttp
from opsdroid.matchers import match_crontab, match_regex
from opsdroid.message import Message
_LOGGER = logging.getLogger(__name__)
################################################################################
# Helper functions #
################################################################################
def human_bytes(num, suffix='B'):
"""Returns a string of human readable file size."""
for unit in ['','Ki','Mi','Gi','Ti','Pi','Ei','Zi']:
if abs(num) < 1024.0:
return "%3.1f%s%s" % (num, unit, suffix)
num /= 1024.0
return "%.1f%s%s" % (num, 'Yi', suffix)
async def count_total_file_size(files):
"""Counts the total size of a list of slack files."""
total_file_size = 0
for file in files:
total_file_size = total_file_size + file["size"]
return total_file_size
async def download_file(slack_api_token, file):
"""Downloads a file from slack into memory and returns the raw bytes."""
headers = {'Authorization': 'Bearer {}'.format(slack_api_token)}
async with aiohttp.ClientSession() as session:
async with session.get(file["url_private"], headers=headers) as resp:
return await resp.read()
async def upload_file(client, file, data, bucket, prefix):
"""Uploads an array of raw bytes to S3 as an object."""
filename = "{}-{}".format(file["id"], file["name"])
resp = await client.put_object(Bucket=bucket,
Key="{}/{}".format(prefix, filename),
Body=data)
if resp["ResponseMetadata"]["HTTPStatusCode"] == 200:
return True
return False
async def cleanup_file(slack_api_token, file):
"""Deletes a file from Slack."""
headers = {"Authorization": "Bearer {}".format(slack_api_token)}
data = {"file": file["id"]}
async with aiohttp.ClientSession() as session:
async with session.post('https://slack.com/api/files.delete', headers=headers, data=data) as resp:
if resp.status == 200:
return True
return False
async def get_file_list(slack_api_token):
"""Gets a list of all files in a slack account."""
all_files = []
total_files = 0
page = 1
pages = None
async with aiohttp.ClientSession() as session:
while pages is None or page <= pages:
async with session.get('https://slack.com/api/files.list?token={}&page={}'.format(slack_api_token, page)) as resp:
if resp.status != 200:
_LOGGER.error("Bad response from slack api: %s", resp.status)
files = await resp.json()
if "error" in files:
if files["error"] == 'user_is_bot':
_LOGGER.error("The Slack token you've used is for a bot user and cannot list files in a team. "
"See the Slack documentation for info on which scope you require https://api.slack.com/methods/files.list")
else:
_LOGGER.error("Unknown error %s", files)
pages = files["paging"]["pages"]
all_files = all_files + files["files"]
total_files = total_files + len(files["files"])
page = page + 1
return all_files
################################################################################
# Skills #
################################################################################
@match_crontab("0 10 * * *")
@match_regex(r'check slack file quota', case_sensitive=False)
async def check_slack_file_quota(opsdroid, config, message):
try:
aws_access_key_id = config["aws_access_key_id"]
aws_secret_access_key = config["aws_secret_access_key"]
slack_api_token = config["slack_api_token"]
s3_region_name = config["s3_region_name"]
max_total_file_size = config["max_total_file_size"]
s3_bucket = config["s3_bucket"]
s3_prefix = config.get("s3_prefix", "")
file_size_buffer = config.get("file_size_buffer", 0)
except KeyError:
_LOGGER.error("Missing config item(s) in skill %s.",
config.get('name', 'aws-tag-compliance'))
return
if message is None:
message = Message("",
None,
config.get("room", connector.default_room),
opsdroid.default_connector)
else:
await message.respond("I'm on it!")
files_removed = 0
data_saved = 0
files = await get_file_list(slack_api_token)
size_threshold = max_total_file_size
while await count_total_file_size(files) > size_threshold:
if size_threshold == max_total_file_size:
size_threshold = max_total_file_size - file_size_buffer
session = aiobotocore.get_session()
async with session.create_client('s3', region_name=s3_region_name,
aws_secret_access_key=aws_secret_access_key,
aws_access_key_id=aws_access_key_id) as client:
data = await download_file(slack_api_token, files[-1])
if await upload_file(client, files[-1], data, s3_bucket, s3_prefix):
if await cleanup_file(slack_api_token, files[-1]):
_LOGGER.debug("Uploaded %s to S3", files[-1]["name"])
files_removed = files_removed + 1
data_saved = data_saved + files[-1]["size"]
files.remove(files[-1])
else:
_LOGGER.debug("%s uploaded to S3 but failed to clean up on Slack", files[-1]["name"])
else:
_LOGGER.debug("Upload of %s failed", files[-1]["name"])
if files_removed > 0:
await message.respond("You were getting close to your Slack file limit so I've moved {} files to the {} bucket on S3 saving {}.".format(files_removed, s3_bucket, human_bytes(data_saved)))
else:
if message.regex:
await message.respond("Nothing to do, file size is {} and quota is {}".format(
human_bytes(await count_total_file_size(files)), human_bytes(max_total_file_size)))