#!/usr/bin/env python3

import argparse
import base64
import json
import os
import re
import socket
import struct
import subprocess
import sys
from pathlib import Path

script_dir = Path(__file__).resolve().parent
for directory in (script_dir.parent / 'lib', script_dir.parent / 'lib' / 'courier'):
    if directory.is_dir():
        sys.path.insert(0, str(directory))

from courier_git import AttachmentError, attach_https, attach_ssh, clone_https, clone_ssh

VERSION = "0.1.0"
PROTOCOL_VERSION = 1
MAX_MESSAGE_SIZE = 24 * 1024 * 1024
MAX_REQUEST_BODY_SIZE = 32 * 1024
MAX_GIT_REQUEST_BODY_SIZE = 1024 * 1024
MAX_GIT_PUSH_BODY_SIZE = 16 * 1024 * 1024
SYSTEM_CONFIG_PATH = Path("/etc/courier/client.json")


class ProtocolError(Exception):
    pass


def socket_path(override=None):
    if override is not None:
        if not override.is_absolute():
            raise ProtocolError("socket path is not absolute")
        return override
    if SYSTEM_CONFIG_PATH.exists():
        try:
            config = json.loads(SYSTEM_CONFIG_PATH.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, UnicodeDecodeError) as error:
            raise ProtocolError("invalid client configuration") from error
        configured_path = config.get("socket_path") if isinstance(config, dict) else None
        if not isinstance(configured_path, str):
            raise ProtocolError("client configuration has no socket path")
        path = Path(configured_path)
        if not path.is_absolute():
            raise ProtocolError("configured socket path is not absolute")
        return path

    runtime_dir = os.environ.get("XDG_RUNTIME_DIR", f"/run/user/{os.getuid()}")
    return Path(runtime_dir) / "courier" / "broker.sock"


def receive_exactly(connection, size):
    chunks = bytearray()
    while len(chunks) < size:
        chunk = connection.recv(size - len(chunks))
        if not chunk:
            raise ProtocolError("connection closed before response was complete")
        chunks.extend(chunk)
    return bytes(chunks)


def receive_message(connection):
    header = receive_exactly(connection, 4)
    message_size = struct.unpack("!I", header)[0]
    if message_size == 0 or message_size > MAX_MESSAGE_SIZE:
        raise ProtocolError("invalid response size")

    payload = receive_exactly(connection, message_size)
    try:
        message = json.loads(payload.decode("utf-8"))
    except (json.JSONDecodeError, UnicodeDecodeError) as error:
        raise ProtocolError("invalid response") from error
    if not isinstance(message, dict):
        raise ProtocolError("invalid response")
    return message


def send_message(connection, message):
    payload = json.dumps(message, separators=(",", ":")).encode("utf-8")
    if not payload or len(payload) > MAX_MESSAGE_SIZE:
        raise ProtocolError("invalid request size")
    connection.sendall(struct.pack("!I", len(payload)) + payload)


def call_broker(request, socket_override=None):
    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
        connection.settimeout(5)
        connection.connect(str(socket_path(socket_override)))
        send_message(connection, request)
        return receive_message(connection)


def broker_version(socket_override=None):
    response = call_broker(
        {"protocol_version": PROTOCOL_VERSION, "command": "get_version"},
        socket_override,
    )

    if response.get("protocol_version") != PROTOCOL_VERSION:
        raise ProtocolError("unsupported protocol version in response")
    if response.get("ok") is not True:
        raise ProtocolError(response.get("error", "broker request failed"))
    if not isinstance(response.get("broker_version"), str):
        raise ProtocolError("broker response has no version")
    return response["broker_version"]


def parse_credential_header(value):
    match = re.fullmatch(
        r"([^:\r\n]+):\s*\$\{([A-Z][A-Z0-9_]*)\}",
        value,
    )
    if match is None:
        raise ProtocolError(
            "header must use the form 'NAME: ${CREDENTIAL_ID}'"
        )
    return {"name": match.group(1).strip(), "credential": match.group(2)}


def make_request(url, header_values, *, method="GET", data_file=None, git_protocol=None):
    if method not in ("GET", "HEAD", "PUT", "POST"):
        raise ProtocolError("unsupported method")
    if (method in ("PUT", "POST")) != (data_file is not None):
        raise ProtocolError("PUT and POST require --data-file; GET and HEAD do not accept a body")
    request = {
        "protocol_version": PROTOCOL_VERSION,
        "command": "https_request",
        "method": method,
        "url": url,
        "headers": [parse_credential_header(value) for value in header_values],
    }
    if git_protocol is not None:
        if git_protocol != 'version=2':
            raise ProtocolError('unsupported Git protocol header')
        request['git_protocol'] = git_protocol
    if data_file is not None:
        maximum = MAX_GIT_REQUEST_BODY_SIZE if method == 'POST' else MAX_REQUEST_BODY_SIZE
        if method == 'POST' and url.split('?', 1)[0].endswith('.git/git-receive-pack'):
            maximum = MAX_GIT_PUSH_BODY_SIZE
        if str(data_file) == '-':
            body = sys.stdin.buffer.read(maximum + 1)
        else:
            with Path(data_file).open("rb") as body_file:
                body = body_file.read(maximum + 1)
        if not 0 < len(body) <= maximum:
            raise ProtocolError(f"request body must contain 1 to {maximum} bytes")
        request["body_base64"] = base64.b64encode(body).decode("ascii")
    return request


def https_request(url, header_values, socket_override=None, *, method="GET",
                  data_file=None, git_protocol=None):
    request = make_request(url, header_values, method=method,
                           data_file=data_file, git_protocol=git_protocol)
    response = call_broker(request, socket_override)
    if response.get("protocol_version") != PROTOCOL_VERSION:
        raise ProtocolError("unsupported protocol version in response")
    if response.get("ok") is not True:
        raise ProtocolError(response.get("error", "broker request failed"))
    if type(response.get("status")) is not int:
        raise ProtocolError("broker response has no HTTP status")
    body_base64 = response.get("body_base64")
    if not isinstance(body_base64, str):
        raise ProtocolError("broker response has no body")
    try:
        body = base64.b64decode(body_base64, validate=True)
    except (ValueError, base64.binascii.Error) as error:
        raise ProtocolError("broker response body is invalid") from error
    return response["status"], body


def stream_https_request(url, header_values, socket_override=None, *, method="GET",
                         data_file=None, git_protocol=None):
    request = make_request(url, header_values, method=method,
                           data_file=data_file, git_protocol=git_protocol)
    request["response_mode"] = "http"
    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
        push = method == 'POST' and url.split('?', 1)[0].endswith('.git/git-receive-pack')
        connection.settimeout(180 if push else 30)
        connection.connect(str(socket_path(socket_override)))
        send_message(connection, request)
        response = receive_message(connection)
        if response.get("protocol_version") != PROTOCOL_VERSION:
            raise ProtocolError("unsupported protocol version in response")
        if response.get("ok") is not True:
            raise ProtocolError(response.get("error", "broker request failed"))
        if response.get("response_mode") != "http":
            raise ProtocolError("broker did not accept HTTP streaming")
        received = False
        while block := connection.recv(65536):
            received = True
            sys.stdout.buffer.write(block)
            sys.stdout.buffer.flush()
        if not received:
            raise ProtocolError("broker returned no HTTP response")


def proxy_service(action):
    if os.geteuid() == 0:
        print('courier-cli: manage the webproxy as your normal user, without sudo', file=sys.stderr)
        return 1

    systemctl = ['/usr/bin/systemctl', '--user']
    unit = 'courier-webproxy.service'
    try:
        if action == 'enable':
            result = subprocess.run(systemctl + ['daemon-reload'], check=False)
            if result.returncode:
                return result.returncode
        arguments = ['status', '--no-pager'] if action == 'status' else [action, '--now']
        return subprocess.run(systemctl + arguments + [unit], check=False).returncode
    except OSError as error:
        print(f'courier-cli: could not manage webproxy user service: {error}', file=sys.stderr)
        return 1


def main(argv=None):
    parser = argparse.ArgumentParser(prog="courier-cli")
    parser.add_argument(
        "--version",
        action="version",
        version=f"%(prog)s {VERSION}",
    )
    parser.add_argument("--socket", type=Path, help=argparse.SUPPRESS)
    subparsers = parser.add_subparsers(dest="command")
    subparsers.add_parser("broker-version", help="show the broker version")
    proxy = subparsers.add_parser('proxy', help='manage webproxy startup for your account, without sudo')
    proxy.add_argument('action', choices=('enable', 'disable', 'status'),
                       help='enable and start, disable and stop, or inspect the user service')
    for transport in ('ssh', 'https'):
        clone = subparsers.add_parser(f'clone-{transport}',
                                     help=f'clone a {transport.upper()} repository through Courier and keep it attached')
        clone.add_argument('url', help=f'{transport.upper()} repository URL')
        clone.add_argument('directory', nargs='?', type=Path, help='destination directory (default: Git chooses from URL)')
        if transport == 'https':
            clone.add_argument('--credential', required=True, help='logical Courier credential ID, never the token')
            clone.add_argument('--port', type=int, default=18080, help='local webproxy port (default: 18080)')
    for command in ('attach-ssh', 'attach-https'):
        attach = subparsers.add_parser(command, help='configure repository-local Git routing without sudo')
        attach.add_argument('repository', nargs='?', type=Path, default=Path('.'))
        attach.add_argument('--replace', action='store_true', help='replace conflicting local routing')
        if command == 'attach-https':
            attach.add_argument('--remote', required=True, help='remote whose stored HTTPS URL should use Courier')
            attach.add_argument('--credential', required=True, help='logical Courier credential ID, never the token')
            attach.add_argument('--port', type=int, default=18080, help='local webproxy port (default: 18080)')
    request_parser = subparsers.add_parser("request", help="perform an authorized HTTPS request")
    request_parser.add_argument("-H", "--header", action="append", required=True)
    request_parser.add_argument("-X", "--method", choices=("GET", "HEAD", "PUT", "POST"), default="GET")
    request_parser.add_argument("--data-file", type=Path, help="request body file; - reads stdin (JSON PUT or Git POST)")
    request_parser.add_argument("--git-protocol", choices=("version=2",), help="Git protocol negotiation header")
    request_parser.add_argument("--stream-http", action="store_true",
                                help="stream a complete HTTP response to stdout")
    request_parser.add_argument("url")
    args = parser.parse_args(argv)

    if args.command == 'proxy':
        return proxy_service(args.action)

    if args.command in ('clone-ssh', 'clone-https', 'attach-ssh', 'attach-https'):
        try:
            if args.command == 'clone-ssh':
                return clone_ssh(args.url, args.directory)
            if args.command == 'clone-https':
                return clone_https(args.url, args.directory, credential=args.credential, port=args.port)
            if args.command == 'attach-ssh':
                attach_ssh(args.repository, replace=args.replace)
            else:
                attach_https(args.repository, remote=args.remote, credential=args.credential,
                             port=args.port, replace=args.replace)
        except (OSError, AttachmentError) as error:
            print(f'courier-cli: {error}', file=sys.stderr)
            return 1
        return 0

    if args.command == "broker-version":
        try:
            version = broker_version(args.socket)
        except (OSError, ProtocolError) as error:
            print(f"courier-cli: could not contact courier-broker: {error}", file=sys.stderr)
            return 1
        print(f"courier-broker {version}")

    if args.command == "request":
        try:
            if args.stream_http:
                stream_https_request(args.url, args.header, args.socket,
                                     method=args.method, data_file=args.data_file, git_protocol=args.git_protocol)
            else:
                _, body = https_request(args.url, args.header, args.socket,
                                        method=args.method, data_file=args.data_file, git_protocol=args.git_protocol)
                sys.stdout.buffer.write(body)
        except (OSError, ProtocolError) as error:
            print(f"courier-cli: broker request failed: {error}", file=sys.stderr)
            return 1

    return 0


if __name__ == "__main__":
    raise SystemExit(main())
