public_skill_preflight.py

"""公开Git托管平台的固定版本Skill根仓库检查。"""
import asyncio
import re
from urllib.parse import quote, urlencode

from proxy_resources import ProxyResources
from skill_archive_source import public_skill_source, read_metadata_response, read_metadata_body, SkillSourceError, METADATA_SECONDS
from skill_manifest import MAX_TREE_ENTRIES, skill_segment, skill_sha, skill_tree

MAX_REQUESTS = 256
PAGE_SIZE = 100


def mirror_head(body):
    if not isinstance(body, bytes):
        raise SkillSourceError('invalid_metadata')
    position = 0
    def packet():
        nonlocal position
        header = body[position:position + 4]
        if len(header) != 4 or not re.fullmatch(b'[0-9a-fA-F]{4}', header):
            raise SkillSourceError('invalid_metadata')
        size = int(header, 16)
        position += 4
        if size == 0:
            return None
        if not 4 < size <= 65520 or position + size - 4 > len(body):
            raise SkillSourceError('invalid_metadata')
        value = body[position:position + size - 4]
        position += size - 4
        return value
    if packet() != b'# service=git-upload-pack\n' or packet() is not None:
        raise SkillSourceError('invalid_metadata')
    head, head_ref, seen = None, None, {}
    for index in range(513):
        value = packet()
        if value is None:
            if head is None or position != len(body) or head_ref is not None and seen.get(head_ref) != head:
                raise SkillSourceError('invalid_metadata')
            return head
        if index == 512 or not value.endswith(b'\n'):
            raise SkillSourceError('invalid_metadata')
        value = value[:-1]
        if index == 0:
            parts = value.split(b'\0')
            if len(parts) != 2 or len(parts[1]) > 16384 or any(c < 32 or c > 126 for c in parts[1]):
                raise SkillSourceError('invalid_metadata')
            capabilities = parts[1].split(b' ')
            if any(cap.startswith(b'object-format=') and cap != b'object-format=sha1' for cap in capabilities):
                raise SkillSourceError('invalid_metadata')
            symbolic = [cap[len(b'symref=HEAD:'):] for cap in capabilities if cap.startswith(b'symref=HEAD:')]
            if len(symbolic) > 1 or symbolic and not symbolic[0].startswith(b'refs/heads/'):
                raise SkillSourceError('invalid_metadata')
            head_ref = symbolic[0] if symbolic else None
            value = parts[0]
        if b'\0' in value or len(value) > 4096:
            raise SkillSourceError('invalid_metadata')
        parts = value.split(b' ')
        if (len(parts) != 2 or any(c < 33 or c > 126 for c in parts[1])
            or not re.fullmatch(b'[0-9a-fA-F]{40}', parts[0]) or parts[0] == b'0' * 40
            or parts[1] in seen):
            raise SkillSourceError('invalid_metadata')
        seen[parts[1]] = parts[0].decode('ascii').lower()
        if index == 0:
            if parts[1] != b'HEAD':
                raise SkillSourceError('invalid_metadata')
            head = parts[0].decode('ascii').lower()
    raise SkillSourceError('invalid_metadata')


def _branch(value):
    try:
        valid = (isinstance(value, str) and value and len(value.encode('utf-8')) <= 4096
            and all(skill_segment(part) for part in value.split('/')))
    except UnicodeError:
        valid = False
    if not valid:
        raise SkillSourceError('invalid_metadata')
    return value


def _normalized_rows(values, provider, path, commit):
    if not isinstance(values, list) or len(values) > MAX_TREE_ENTRIES:
        raise SkillSourceError('invalid_metadata')
    result, aliases = {}, set()
    for entry in values:
        if not isinstance(entry, dict) or not isinstance(entry.get('path'), str):
            raise SkillSourceError('invalid_metadata')
        prefix = path + '/' if path else ''
        if not entry['path'].startswith(prefix):
            raise SkillSourceError('invalid_metadata')
        name = entry['path'][len(prefix):]
        if not skill_segment(name) or name.casefold() in aliases:
            raise SkillSourceError('invalid_metadata')
        aliases.add(name.casefold())
        if provider == 'gitlab.com':
            if entry.get('name') != name:
                raise SkillSourceError('invalid_metadata')
            sha = skill_sha(entry.get('id'))
            kind, mode = entry.get('type'), entry.get('mode')
            if (kind, mode) not in {('blob', '100644'), ('blob', '100755'), ('blob', '120000'), ('tree', '040000'), ('commit', '160000')}:
                raise SkillSourceError('invalid_metadata')
        else:
            info = entry.get('commit')
            if not isinstance(info, dict) or info.get('type') != 'commit' or skill_sha(info.get('hash')) != commit:
                raise SkillSourceError('invalid_metadata')
            attributes = entry.get('attributes', [] if entry.get('type') == 'commit_directory' else None)
            if not isinstance(attributes, list) or len(attributes) > 32 or any(not isinstance(v, str) or len(v) > 128 for v in attributes):
                raise SkillSourceError('invalid_metadata')
            if len(set(attributes)) != len(attributes) or not set(attributes) <= {'executable', 'link', 'subrepository', 'binary'}:
                raise SkillSourceError('invalid_metadata')
            kind = entry.get('type')
            if kind == 'commit_directory':
                kind, mode = 'tree', '040000'
                if attributes:
                    raise SkillSourceError('invalid_metadata')
            elif kind == 'commit_file':
                if type(entry.get('size')) is not int or entry['size'] < 0:
                    raise SkillSourceError('invalid_metadata')
                kind = 'blob'
                if 'subrepository' in attributes:
                    if len(attributes) != 1:
                        raise SkillSourceError('invalid_metadata')
                    kind, mode = 'commit', '160000'
                else:
                    mode = '120000' if 'link' in attributes else '100755' if 'executable' in attributes else '100644'
            else:
                raise SkillSourceError('invalid_metadata')
            sha = commit
        result[name] = dict(path=name, sha=sha, type=kind, mode=mode)
    return result


async def probe_public_skill(item, *, send, client_factory, transport_factory):
    source = public_skill_source(item)
    host, owner, repo = source['host'], source['owner'], source['repo']
    client = client_factory(follow_redirects=False, trust_env=False, transport=transport_factory())
    resources = ProxyResources(client)
    headers = {'Accept': 'application/json', 'Accept-Encoding': 'identity', 'User-Agent': 'ComfyUI-Ranking-SaaS'}
    requests = 0
    async def response(url, accept='application/json'):
        nonlocal requests
        requests += 1
        if requests > MAX_REQUESTS:
            raise SkillSourceError('metadata_too_large')
        return resources.track(await send(client, 'GET', url, headers=headers | {'Accept': accept}, timeout=METADATA_SECONDS, max_redirects=0))
    async def read(url, paginated=False):
        return await read_metadata_response(await response(url), resources, allow_pagination=paginated)

    try:
        async with asyncio.timeout(METADATA_SECONDS):
            if host == 'kkgithub.com':
                advertisement = await response(f'https://kkgithub.com/{owner}/{repo}.git/info/refs?service=git-upload-pack', 'application/x-git-upload-pack-advertisement')
                content_type = advertisement.headers.get('content-type', '').split(';', 1)[0].strip().lower()
                body = await read_metadata_body(advertisement, resources)
                if content_type != 'application/x-git-upload-pack-advertisement':
                    raise SkillSourceError('invalid_metadata')
                commit = mirror_head(body)
                base = f'https://api.github.com/repos/{owner}/{repo}'
                metadata = await read(base + '/git/commits/' + commit)
                if not isinstance(metadata, dict) or skill_sha(metadata.get('sha')) != commit or not isinstance(metadata.get('tree'), dict):
                    raise SkillSourceError('invalid_metadata')
                root_tree = skill_sha(metadata['tree'].get('sha'))
            else:
                if host == 'gitlab.com':
                    base = 'https://gitlab.com/api/v4/projects/' + quote(owner + '/' + repo, safe='')
                    repository = await read(base)
                    branch = _branch(repository.get('default_branch') if isinstance(repository, dict) else None)
                    metadata = await read(base + '/repository/branches/' + quote(branch, safe=''))
                    if not isinstance(metadata, dict) or metadata.get('name') != branch or not isinstance(metadata.get('commit'), dict):
                        raise SkillSourceError('invalid_metadata')
                    commit = skill_sha(metadata['commit'].get('id'))
                elif host == 'gitee.com':
                    base = f'https://gitee.com/api/v5/repos/{owner}/{repo}'
                    repository = await read(base)
                    branch = _branch(repository.get('default_branch') if isinstance(repository, dict) else None)
                    metadata = await read(base + '/branches/' + quote(branch, safe=''))
                    if not isinstance(metadata, dict) or metadata.get('name') != branch or not isinstance(metadata.get('commit'), dict):
                        raise SkillSourceError('invalid_metadata')
                    commit = skill_sha(metadata['commit'].get('sha'))
                    details = metadata['commit'].get('commit')
                    if not isinstance(details, dict) or not isinstance(details.get('tree'), dict):
                        raise SkillSourceError('invalid_metadata')
                    root_tree = skill_sha(details['tree'].get('sha'))
                else:
                    base = f'https://api.bitbucket.org/2.0/repositories/{owner}/{repo}'
                    repository = await read(base)
                    mainbranch = repository.get('mainbranch') if isinstance(repository, dict) else None
                    branch = _branch(mainbranch.get('name') if isinstance(mainbranch, dict) else None)
                    metadata = await read(base + '/refs/branches/' + quote(branch, safe=''))
                    if not isinstance(metadata, dict) or metadata.get('name') != branch or metadata.get('type') != 'branch' or not isinstance(metadata.get('target'), dict):
                        raise SkillSourceError('invalid_metadata')
                    if metadata['target'].get('type') != 'commit':
                        raise SkillSourceError('invalid_metadata')
                    commit = skill_sha(metadata['target'].get('hash'))

            async def directory(path='', tree_id=None):
                if host in ('gitee.com', 'kkgithub.com'):
                    value = await read(base + '/git/trees/' + (tree_id or root_tree))
                    return skill_tree(value, tree_id or root_tree)
                values = []
                for page in range(1, MAX_TREE_ENTRIES + 2):
                    if host == 'gitlab.com':
                        url = base + '/repository/tree?' + urlencode(dict(ref=commit, path=path, per_page=PAGE_SIZE, page=page))
                        value = await read(url, True)
                        if not isinstance(value, list) or len(value) > PAGE_SIZE:
                            raise SkillSourceError('invalid_metadata')
                        more = bool(value)
                        entries = value
                    else:
                        url = base + '/src/' + commit + '/' + (quote(path, safe='') + '/' if path else '') + '?' + urlencode(dict(pagelen=PAGE_SIZE, page=page))
                        value = await read(url, True)
                        if (not isinstance(value, dict) or type(value.get('page')) is not int or value['page'] != page
                            or type(value.get('pagelen')) is not int or not 1 <= value['pagelen'] <= PAGE_SIZE
                            or not isinstance(value.get('values'), list) or len(value['values']) > value['pagelen']):
                            raise SkillSourceError('invalid_metadata')
                        next_page = value.get('next')
                        if next_page is not None and (not isinstance(next_page, str) or not 0 < len(next_page) <= 8192):
                            raise SkillSourceError('invalid_metadata')
                        more, entries = next_page is not None, value['values']
                        if more and not entries:
                            raise SkillSourceError('invalid_metadata')
                    values.extend(entries)
                    if len(values) > MAX_TREE_ENTRIES:
                        raise SkillSourceError('metadata_too_large')
                    if not more:
                        return _normalized_rows(values, host, path, commit)
                raise SkillSourceError('metadata_too_large')

            await directory()
            return True
    finally:
        await resources.close()