#!/usr/bin/env python3
"""One explicit retrieval. Repeated IP searches consume quota on every retrieval."""
import argparse
import json
import os
import tempfile
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--url', required=True, help='STIX API URL, including IP or since/until parameters')
    parser.add_argument('--output', default='senda.stix.json')
    args = parser.parse_args()
    url = urllib.parse.urlsplit(args.url)
    if url.scheme != 'https' or not url.netloc or url.username or url.password:
        parser.error('Use an HTTPS URL without embedded credentials')
    key = os.environ.get('SENDA_API_KEY', '').strip()
    if not key:
        parser.error('Set SENDA_API_KEY in the environment')
    request = urllib.request.Request(args.url, headers={'Authorization': 'Bearer ' + key,
                                                        'Accept': 'application/stix+json'})
    try:
        with urllib.request.build_opener(NoRedirect()).open(request, timeout=900) as response:
            document = json.load(response)
            headers = response.headers
        if not isinstance(document, dict) or document.get('type') != 'bundle' or not isinstance(document.get('objects'), list):
            raise ValueError('Response is not a STIX bundle')
        path = Path(args.output).resolve()
        path.parent.mkdir(parents=True, exist_ok=True)
        temporary = None
        try:
            with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8', dir=path.parent, delete=False) as stream:
                temporary = Path(stream.name)
                json.dump(document, stream, ensure_ascii=False, indent=2)
                stream.flush()
                os.fsync(stream.fileno())
            os.replace(temporary, path)
        finally:
            if temporary and temporary.exists():
                temporary.unlink()
        print('Saved', path, 'objects=' + str(len(document['objects'])),
              'searches=' + headers.get('X-Senda-Search-Count', '?'),
              'remaining=' + headers.get('X-Senda-API-Remaining', '?'))
        cursor = headers.get('X-Senda-Next-Cursor')
        if cursor:
            print('Next cursor:', cursor)
            print('Request the same observation window with &cursor=... to retrieve the next page.')
    except urllib.error.HTTPError as exc:
        print('Retrieval failed:', exc.code, 'Retry-After:', exc.headers.get('Retry-After', '-'))
        raise SystemExit(1)
    except (OSError, ValueError, urllib.error.URLError):
        print('Retrieval failed; the previous output file was retained.')
        raise SystemExit(1)


if __name__ == '__main__':
    main()
