-
Notifications
You must be signed in to change notification settings - Fork 0
/
image_storage.py
217 lines (178 loc) · 6.47 KB
/
image_storage.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
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
# vim: ts=4 et:
import copy
import hashlib
import shutil
import os
from glob import glob
from pathlib import Path
from subprocess import Popen, PIPE
from urllib.parse import urlparse
from image_tags import DictObj
def run(cmd, log, errmsg=None, errvals=[], err_ok=False):
# ensure command and error values are lists of strings
cmd = [str(c) for c in cmd]
errvals = [str(ev) for ev in errvals]
log.debug('COMMAND: %s', ' '.join(cmd))
p = Popen(cmd, stdout=PIPE, stdin=PIPE, stderr=PIPE, encoding='utf8')
out, err = p.communicate()
if p.returncode:
if errmsg:
if err_ok:
log.debug(errmsg, *errvals)
else:
log.error(errmsg, *errvals)
log.debug('EXIT: %d / COMMAND: %s', p.returncode, ' '.join(cmd))
log.debug('STDOUT:\n%s', out)
log.debug('STDERR:\n%s', err)
raise RuntimeError
return out, err
class ImageStorage():
def __init__(self, local, storage_url, log):
self.log = log
self.local = local
self.url = storage_url.removesuffix('/')
url = urlparse(self.url)
if url.scheme not in ['', 'file', 'ssh']:
self.log.error('Storage with "%s" scheme is unsupported', url.scheme)
raise RuntimeError
if url.scheme in ['', 'file']:
self.scheme = 'file'
self.remote = Path(url.netloc + url.path).expanduser()
else:
self.scheme = 'ssh'
self.host = url.hostname
self.remote = Path(url.path[1:]) # drop leading / -- use // for absolute path
self.ssh = DictObj({
'port': ['-p', url.port] if url.port else [],
'user': ['-l', url.username] if url.username else [],
})
self.scp = DictObj({
'port': ['-P', url.port] if url.port else [],
'user': url.username + '@' if url.username else '',
})
def _checksum(self, file, save=False):
log = self.log
src = self.local
log.debug("Calculating checksum for '%s'", file)
sha512_hash = hashlib.sha512()
with open(src / file, 'rb') as f:
for block in iter(lambda: f.read(4096), b''):
sha512_hash.update(block)
checksum = sha512_hash.hexdigest()
if save:
log.debug("Saving '%s'", file + '.sha512')
with open(str(src / file) + '.sha512', 'w') as f:
print(checksum, file=f)
return checksum
def store(self, *files, checksum=False):
log = self.log
src = self.local
dest = self.remote
# take care of any globbing in file list
files = [Path(p).name for p in sum([glob(str(src / f)) for f in files], [])]
if not files:
log.debug('No files to store')
return
if checksum:
log.info('Creating checksum(s) for %s', files)
for f in copy.copy(files):
self._checksum(f, save=True)
files.append(f + '.sha512')
log.info('Storing %s', files)
if self.scheme == 'file':
dest.mkdir(parents=True, exist_ok=True)
for file in files:
shutil.copy2(src / file, dest / file)
return
url = self.url
host = self.host
ssh = self.ssh
scp = self.scp
run(
['ssh'] + ssh.port + ssh.user + [host, 'mkdir', '-p', dest],
log=log, errmsg='Unable to ensure existence of %s', errvals=[url]
)
src_files = []
for file in files:
log.info('Storing %s', url + '/' + file)
src_files.append(src / file)
run(
['scp'] + scp.port + src_files + [scp.user + ':'.join([host, str(dest)])],
log=log, errmsg='Failed to store files'
)
def retrieve(self, *files, checksum=False):
log = self.log
# TODO? use list()
if not files:
log.debug('No files to retrieve')
return
src = self.remote
dest = self.local
dest.mkdir(parents=True, exist_ok=True)
if self.scheme == 'file':
for file in files:
log.debug('Retrieving %s', src / file)
shutil.copy2(src / file, dest / file)
return
url = self.url
host = self.host
scp = self.scp
src_files = []
for file in files:
log.debug('Retrieving %s', url + '/' + file)
src_files.append(scp.user + ':'.join([host, str(src / file)]))
run(
['scp'] + scp.port + src_files + [dest],
log=log, errmsg='Failed to retrieve files'
)
# TODO: optional files=[]?
def list(self, match=None, err_ok=False):
log = self.log
path = self.remote
if not match:
match = '*'
files = []
if self.scheme == 'file':
path.mkdir(parents=True, exist_ok=True)
log.debug('Listing of %s files in %s', match, path)
files = sorted(glob(str(path / match)), key=os.path.getmtime, reverse=True)
else:
url = self.url
host = self.host
ssh = self.ssh
log.debug('Listing %s files at %s', match, url)
run(
['ssh'] + ssh.port + ssh.user + [host, 'mkdir', '-p', path],
log=log, errmsg='Unable to create path', err_ok=err_ok
)
out, _ = run(
['ssh'] + ssh.port + ssh.user + [host, 'ls', '-1drt', path / match],
log=log, errmsg='Failed to list files', err_ok=err_ok
)
files = out.splitlines()
return [os.path.basename(f) for f in files]
def remove(self, *files):
log = self.log
# TODO? use list()
if not files:
log.debug('No files to remove')
return
dest = self.remote
if self.scheme == 'file':
for file in files:
path = dest / file
log.debug('Removing %s', path)
if path.exists():
path.unlink()
return
url = self.url
host = self.host
ssh = self.ssh
dest_files = []
for file in files:
log.debug('Removing %s', url + '/' + file)
dest_files.append(dest / file)
run(
['ssh'] + ssh.port + ssh.user + [host, 'rm', '-f'] + dest_files,
log=log, errmsg='Failed to remove files'
)