feature-test-1 / server.py
rphr's picture
Create server.py
54b7647 verified
Raw
History Blame Contribute Delete
4.75 kB
import asyncio
import json
import socket
import dns.resolver # To handle DNS resolution using custom DNS server
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
import uvicorn
app = FastAPI()
# Function to resolve DNS with a custom DNS server (4.2.2.4)
def resolve_with_custom_dns(domain: str, dns_server: str = "4.2.2.4"):
try:
resolver = dns.resolver.Resolver()
resolver.nameservers = [dns_server] # Set the custom DNS server
answers = resolver.resolve(domain)
for answer in answers:
print(f"Resolved {domain} to {answer.to_text()}")
return [answer.to_text() for answer in answers]
except Exception as e:
print(f"[!] Error resolving {domain}: {e}")
return None
@app.get("/")
async def root():
html_content = """
<!DOCTYPE html>
<html>
<head>
<title>WebSocket Proxy Server</title>
<style>
body {
font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
background-color: #f4f4f9;
color: #333;
display: flex;
justify-content: center;
align-items: center;
height: 100vh;
margin: 0;
}
.container {
background: white;
padding: 40px;
border-radius: 10px;
box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
text-align: center;
}
h1 {
color: #2c3e50;
}
p {
font-size: 1.1em;
color: #555;
}
.status {
display: inline-block;
margin-top: 20px;
padding: 10px 20px;
background-color: #2ecc71;
color: white;
border-radius: 5px;
font-weight: bold;
}
</style>
</head>
<body>
<div class="container">
<h1>Proxy Server Node</h1>
<p>This server is running and ready to accept WebSocket tunnel connections.</p>
<div class="status">System Online</div>
</div>
</body>
</html>
"""
return HTMLResponse(content=html_content, status_code=200)
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
try:
# 1. Receive the target destination from the proxy
init_data = await websocket.receive_text()
target = json.loads(init_data)
host = target.get("host")
port = target.get("port")
# 2. Resolve DNS using custom DNS server (4.2.2.4)
#print(f"[*] Resolving DNS for {host} using 4.2.2.4")
resolved_ips = resolve_with_custom_dns(host, dns_server="4.2.2.4")
if resolved_ips is None:
await websocket.send_text(json.dumps({"status": "error", "message": "DNS resolution failed"}))
await websocket.close()
return
# Pick the first resolved IP address
ip = resolved_ips[0]
#print(f"[*] Resolved {host} to {ip}")
# 3. Open TCP connection to the resolved IP
#print(f"[*] Connecting to remote {ip}:{port}")
remote_reader, remote_writer = await asyncio.open_connection(ip, port)
await websocket.send_text(json.dumps({"status": "success"}))
# 4. Forwarding Loops
async def ws_to_tcp():
try:
while True:
data = await websocket.receive_bytes()
remote_writer.write(data)
await remote_writer.drain()
except Exception:
pass
finally:
remote_writer.close()
async def tcp_to_ws():
try:
while True:
data = await remote_reader.read(4096)
if not data:
break
await websocket.send_bytes(data)
except Exception:
pass
finally:
await websocket.close()
await asyncio.gather(ws_to_tcp(), tcp_to_ws())
except Exception as e:
print(f"[!] Server Error: {e}")
try:
await websocket.send_text(json.dumps({"status": "error", "message": str(e)}))
await websocket.close()
except:
pass
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)