import argparse
import os
import requests
API_BASE_URL = "https://www.corgea.app/api/v1/start-scan"
CHUNK_SIZE = 5 * 1024 * 1024 # 5 MB example; the CLI uses 50 MB. Any size is valid.
TOKEN_HEADERS = {"CORGEA-TOKEN": "<YOUR_TOKEN>"}
def transfer_url(transfer_id):
"""HEAD and PATCH continue URLs require a trailing slash."""
return f"{API_BASE_URL}/{transfer_id}/"
def require_upload_offset(headers):
"""Return the server's next-byte offset. Fail if the header is missing."""
value = headers.get("Upload-Offset")
if value is None:
raise RuntimeError("Response is missing the required Upload-Offset header.")
return int(value)
def initiate_upload(file_path):
"""Start the transfer with the file name and an empty first part."""
file_name = os.path.basename(file_path)
response = requests.post(
API_BASE_URL,
data={"scan_type": "blast"},
files={"files": (file_name, b"")},
headers=TOKEN_HEADERS,
)
response.raise_for_status()
transfer_id = response.json().get("transfer_id")
if not transfer_id:
raise RuntimeError("Start-scan response is missing transfer_id.")
print(f"Upload initiated. Transfer ID: {transfer_id}")
return transfer_id
def check_upload_status(transfer_id):
"""Return the next byte offset from the Upload-Offset header."""
response = requests.head(
transfer_url(transfer_id),
headers=TOKEN_HEADERS,
)
response.raise_for_status()
return require_upload_offset(response.headers)
def upload_chunk(
file_path,
transfer_id,
chunk_offset,
chunk_data,
project_name,
branch,
repo_url,
sha,
partial_scan,
files_to_scan,
metadata="",
):
headers = {
"Upload-Offset": str(chunk_offset),
"Upload-Length": str(os.path.getsize(file_path)),
"Upload-Name": os.path.basename(file_path),
**TOKEN_HEADERS,
}
form_data = {}
if project_name:
form_data["project_name"] = project_name
if branch:
form_data["branch"] = branch
if repo_url:
form_data["repo_url"] = repo_url
if sha:
form_data["sha"] = sha
if partial_scan:
form_data["partial_scan"] = "true"
if files_to_scan:
form_data["files_to_scan"] = files_to_scan
if metadata:
form_data["metadata"] = metadata
response = requests.patch(
transfer_url(transfer_id),
headers=headers,
files={"chunk_data": ("chunk", chunk_data, "application/octet-stream")},
data=form_data,
)
response.raise_for_status()
return response.json(), response.headers
def upload_file_in_chunks(
file_path,
project_name,
branch,
repo_url,
sha,
partial_scan=False,
files_to_scan="",
transfer_id=None,
metadata="",
):
file_size = os.path.getsize(file_path)
if transfer_id:
print(f"Resuming transfer {transfer_id}")
else:
transfer_id = initiate_upload(file_path)
offset = check_upload_status(transfer_id)
print(f"Uploading {file_path} ({file_size} bytes) from offset {offset}")
scan_id = None
with open(file_path, "rb") as handle:
while offset < file_size:
handle.seek(offset)
chunk_data = handle.read(CHUNK_SIZE)
if not chunk_data:
break
payload, response_headers = upload_chunk(
file_path,
transfer_id,
offset,
chunk_data,
project_name,
branch,
repo_url,
sha,
partial_scan,
files_to_scan,
metadata,
)
scan_id = payload.get("scan_id")
if scan_id:
break
message = payload.get("message") or ""
if (
payload.get("status") == "ok"
and "already in progress" in message.lower()
):
print(message)
return
next_offset = require_upload_offset(response_headers)
if next_offset <= offset:
raise RuntimeError(
f"Upload-Offset did not advance (was {offset}, now {next_offset})."
)
if next_offset > file_size:
raise RuntimeError(
f"Upload-Offset {next_offset} exceeds file size {file_size}."
)
offset = next_offset
print(f"Uploaded to server offset {offset} / {file_size} bytes")
if not scan_id:
raise RuntimeError(
"Upload finished without a scan_id. HEAD /start-scan/{transfer_id}/ does not return the scan_id after a completed transfer; look it up with the scans API."
)
print(f"Upload complete. Scan ID: {scan_id}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Upload a BLAST archive in chunks.")
parser.add_argument("file", help="Path to the archive to upload")
parser.add_argument("--branch", required=False)
parser.add_argument("--repo_url", required=False)
parser.add_argument("--sha", required=False)
parser.add_argument("--project_name", required=False)
parser.add_argument("--partial_scan", action="store_true")
parser.add_argument(
"--files_to_scan",
required=False,
default="",
help="Comma-separated files for a partial scan",
)
parser.add_argument(
"--transfer_id",
required=False,
help="Resume an existing transfer instead of starting a new one",
)
parser.add_argument(
"--metadata",
required=False,
default="",
help='JSON object string sent with every chunk, e.g. {"pipeline_url":"https://ci.example/run/123"}',
)
args = parser.parse_args()
if not os.path.isfile(args.file):
raise SystemExit(f"File '{args.file}' does not exist.")
upload_file_in_chunks(
args.file,
args.project_name,
args.branch,
args.repo_url,
args.sha,
args.partial_scan,
args.files_to_scan,
args.transfer_id,
args.metadata,
)
