1 | import hashlib |
---|
2 | import logging |
---|
3 | import os |
---|
4 | import shlex |
---|
5 | import shutil |
---|
6 | import subprocess |
---|
7 | import urllib.request |
---|
8 | |
---|
9 | def _compute_md5(path, bs=2**20): |
---|
10 | m = hashlib.md5() |
---|
11 | with open(path, 'rb') as f: |
---|
12 | while True: |
---|
13 | buf = f.read(bs) |
---|
14 | if not buf: |
---|
15 | break |
---|
16 | m.update(buf) |
---|
17 | return m.hexdigest() |
---|
18 | |
---|
19 | |
---|
20 | def tip_fetch_csum(tip_addr, image_name): |
---|
21 | """ |
---|
22 | """ |
---|
23 | url = f'http://{tip_addr}:9999/{image_name}.img.full.sum' |
---|
24 | try: |
---|
25 | with urllib.request.urlopen(f'{url}') as resp: |
---|
26 | r = resp.readline().rstrip().decode('utf-8') |
---|
27 | except urllib.error.URLError as e: |
---|
28 | logging.warning('URL error when fetching checksum: {e.reason}') |
---|
29 | raise e |
---|
30 | except urllib.error.HTTPError as e: |
---|
31 | logging.warning(f'HTTP Error when fetching checksum: {e.reason}') |
---|
32 | raise e |
---|
33 | return r |
---|
34 | |
---|
35 | |
---|
36 | def tip_write_csum(image_name): |
---|
37 | """ |
---|
38 | TODO: Check for CACHE partition |
---|
39 | """ |
---|
40 | image_path = f'/opt/opengnsys/cache/opt/opengnsys/images/{image_name}.img' |
---|
41 | |
---|
42 | if not os.path.exists(image_path): |
---|
43 | logging.error('Invalid image path') |
---|
44 | raise ValueError('Invalid image path for tiptorrent checksum writing') |
---|
45 | |
---|
46 | filename = image_path + ".full.sum" |
---|
47 | csum = _compute_md5(image_path) |
---|
48 | with open(filename, 'w') as f: |
---|
49 | f.write(csum) |
---|
50 | |
---|
51 | return csum |
---|
52 | |
---|
53 | |
---|
54 | def tip_check_csum(tip_addr, image_name): |
---|
55 | """ |
---|
56 | """ |
---|
57 | image_path = f'/opt/opengnsys/cache/opt/opengnsys/images/{image_name}.img' |
---|
58 | if not os.path.exists(image_path): |
---|
59 | logging.error('Invalid image path') |
---|
60 | raise ValueError('Invalid image path for tiptorrent image csum comparison') |
---|
61 | |
---|
62 | cache_csum = _compute_md5(image_path) |
---|
63 | remote_csum = tip_fetch_csum(tip_addr, image_name) |
---|
64 | logging.debug(f'cache_csum: {cache_csum}') |
---|
65 | logging.debug(f'remote_csum: {remote_csum}') |
---|
66 | return cache_csum == remote_csum |
---|
67 | |
---|
68 | |
---|
69 | def tip_client_get(tip_addr, image_name): |
---|
70 | """ |
---|
71 | """ |
---|
72 | logging.info(f'Fetching image {image_name} from tiptorrent server at {tip_addr}') |
---|
73 | cmd = f'tiptorrent-client {tip_addr} {image_name}.img' |
---|
74 | logfile = open('/tmp/command.log', 'wb', 0) |
---|
75 | |
---|
76 | try: |
---|
77 | proc = subprocess.Popen(shlex.split(cmd), |
---|
78 | stdout=logfile, |
---|
79 | cwd='/opt/opengnsys/cache/opt/opengnsys/images/') |
---|
80 | proc.communicate() |
---|
81 | except: |
---|
82 | logging.error('Exception when running tiptorrent client GET subprocess') |
---|
83 | raise ValueError('Unexpected error running tiptorrent subprocess') |
---|
84 | finally: |
---|
85 | logfile.close() |
---|
86 | |
---|
87 | if proc.returncode != 0: |
---|
88 | logging.error(f'Error fetching image {image_name} via tiptorrent') |
---|
89 | raise ValueError('Tiptorrent download failed') |
---|
90 | else: |
---|
91 | logging.info('tiptorrent transfer completed, writing checksum') |
---|
92 | tip_write_csum(image_name) |
---|