File size: 3,670 Bytes
9540b54
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0b0a9fa
 
9540b54
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import uuid
from pydantic import BaseModel
import asyncio
from typing import Dict, Optional

app = FastAPI()

# Enable CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Store connections and their states
class Connection:
    def __init__(self):
        self.websocket: Optional[WebSocket] = None
        self.offer: Optional[Dict] = None
        self.answer: Optional[Dict] = None

connections: Dict[str, Connection] = {}

class PeerID(BaseModel):
    peer_id: str

class WebRTCOffer(BaseModel):
    peer_id: str
    offer: Dict

class WebRTCAnswer(BaseModel):
    peer_id: str
    answer: Dict

@app.websocket("/ws/{peer_id}")
async def websocket_endpoint(websocket: WebSocket, peer_id: str):
    await websocket.accept()
    
    if peer_id not in connections:
        connections[peer_id] = Connection()
    
    connections[peer_id].websocket = websocket
    
    try:
        while True:
            data = await websocket.receive_json()
            
            if data["type"] == "offer":
                connections[peer_id].offer = data["offer"]
                if connections[peer_id].answer:
                    await websocket.send_json({
                        "type": "answer",
                        "answer": connections[peer_id].answer
                    })
                    
            elif data["type"] == "answer":
                connections[peer_id].answer = data["answer"]
                if connections[peer_id].websocket:
                    await connections[peer_id].websocket.send_json({
                        "type": "answer",
                        "answer": data["answer"]
                    })
                    
    except WebSocketDisconnect:
        if peer_id in connections:
            del connections[peer_id]

@app.post("/register_peer")
async def register_peer(peer: PeerID):
    if peer.peer_id not in connections:
        connections[peer.peer_id] = Connection()
    return JSONResponse({
        "status": "registered",
        "peer_id": peer.peer_id
    })

@app.post("/send_offer")
async def send_offer(offer: WebRTCOffer):
    if offer.peer_id in connections:
        connections[offer.peer_id].offer = offer.offer
        return JSONResponse({"status": "offer_stored"})
    return JSONResponse(
        status_code=404,
        content={"status": "error", "message": "Peer not found"}
    )

@app.get("/get_offer/{peer_id}")
async def get_offer(peer_id: str):
    if (peer_id in connections and 
        connections[peer_id].offer):
        return JSONResponse({
            "offer": connections[peer_id].offer
        })
    return JSONResponse({"status": "no_offer"})

@app.post("/send_answer")
async def send_answer(answer: WebRTCAnswer):
    if answer.peer_id in connections:
        connections[answer.peer_id].answer = answer.answer
        return JSONResponse({"status": "answer_stored"})
    return JSONResponse(
        status_code=404,
        content={"status": "error", "message": "Peer not found"}
    )

@app.get("/get_answer/{peer_id}")
async def get_answer(peer_id: str):
    if (peer_id in connections and 
        connections[peer_id].answer):
        answer = connections[peer_id].answer
        connections[peer_id].answer = None
        return JSONResponse({"answer": answer})
    return JSONResponse({"status": "no_answer"})

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)