hail/devbin/trivy-scan-trigger.py at 4ca7529724c6ca8a9b271ca6cda2d058a64506f0 · cjllanwarne/hail · 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
#!/usr/bin/env python3
"""Find the most recent completed ci/deploy/prod batch and extract the batch image name."""
import itertools
import json
import re
import subprocess
import sys
import threading
import time


def run(args):
    result = subprocess.run(args, capture_output=True, text=True)
    if result.returncode != 0:
        print(result.stderr, file=sys.stderr)
        sys.exit(1)
    return result.stdout


def print_placeholder_during_lookup(name, getter):
    stop = threading.Event()
    result = [None]
    prefix = f'{name} : '

    def spin():
        for frame in itertools.cycle('⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏'):
            if stop.is_set():
                break
            print(f'\r{prefix}{frame}', end='', flush=True)
            time.sleep(0.1)

    t = threading.Thread(target=spin, daemon=True)
    t.start()
    try:
        result[0] = getter()
    finally:
        stop.set()
        t.join()

    line = f'{prefix}{result[0]}'
    print(f'\r{line:<60}')
    return result[0]


def find_batch_id():
    out = run(['hailctl', 'batch', 'list', '--query', 'batch_type=ci/deploy/prod complete', '--limit', '1', '-o', 'json'])
    if not out.strip() or not out.strip().startswith('['):
        print('No completed ci/deploy/prod batches found', file=sys.stderr)
        sys.exit(1)
    batches = json.loads(out)
    if not batches:
        print('No completed ci/deploy/prod batches found', file=sys.stderr)
        sys.exit(1)
    return batches[0]['id']


def find_sha(batch_id):
    batch = json.loads(run(['hailctl', 'batch', 'get', str(batch_id), '-o', 'json']))[0]
    return batch['attributes']['sha']


def find_image(batch_id):
    out = run(['hailctl', 'batch', 'jobs', str(batch_id), '--name', 'batch_image', '-o', 'json'])
    if not out.strip() or not out.strip().startswith('['):
        print(f'No batch_image job found in batch {batch_id}', file=sys.stderr)
        sys.exit(1)
    jobs = json.loads(out)
    if not jobs:
        print(f'No batch_image job found in batch {batch_id}', file=sys.stderr)
        sys.exit(1)
    job_id = jobs[0]['job_id']
    job = json.loads(run(['hailctl', 'batch', 'job', str(batch_id), str(job_id), '-o', 'json']))[0]
    command_parts = job['spec']['process']['command']
    command = next((p for p in command_parts if '--output' in p), None)
    if command is None:
        print('No --output flag found in batch_image job command', file=sys.stderr)
        sys.exit(1)
    match = re.search(r"--output 'type=image,\"name=([^,\"]+)", command)
    if not match:
        print('Could not find image name in batch_image job command', file=sys.stderr)
        sys.exit(1)
    return match.group(1)


def main():
    print('Searching latest completed deploy batch for commit sha and image to scan...')
    print()
    batch_id = print_placeholder_during_lookup('  Batch', find_batch_id)
    sha      = print_placeholder_during_lookup('  SHA  ', lambda: find_sha(batch_id))
    image    = print_placeholder_during_lookup('  Image', lambda: find_image(batch_id))

    BRANCH = 'main'
    gh_cmd = (
        f'gh workflow run trivy-scan.yml --repo hail-is/hail --ref {BRANCH}'
        f' -f branch={BRANCH} -f commit_hash={sha} -f images={image}'
    )
    gh_cmd_display = (
        f'gh workflow run trivy-scan.yml \\\n'
        f'    --repo hail-is/hail --ref {BRANCH} \\\n'
        f'    -f branch={BRANCH} \\\n'
        f'    -f commit_hash={sha} \\\n'
        f'    -f images={image}'
    )

    print()
    print(f'Proposed command:\n{gh_cmd_display}')
    print()
    try:
        answer = input('Run it now? [y/N] ').strip().lower()
    except EOFError:
        answer = ''
    if answer != 'y':
        print('Aborted.')
        sys.exit(0)

    subprocess.run(
        [
            'gh', 'workflow', 'run', 'trivy-scan.yml',
            '--repo', 'hail-is/hail',
            '--ref', BRANCH,
            '-f', f'branch={BRANCH}',
            '-f', f'commit_hash={sha}',
            '-f', f'images={image}',
        ],
        check=True,
    )
    print('Workflow triggered.')


if __name__ == '__main__':
    main()