p2s.py

from random import randint
import os
import platform
import socket
import traceback
import argparse
from rfc import RFC

class PeerToServer: 
    def __init__(self, peer_ip, peer_port, server_ip, server_port=7734):
        self.peer_ip = peer_ip
        self.peer_port = peer_port if peer_port else randint(1024, 65535)
        self.server_ip = server_ip
        self.server_port = server_port
        self.server_connection = None

    def get_connection_to_server(self):
        self.server_connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_connection.connect((self.server_ip, self.server_port))

    def send_add_request(self, rfc_number, rfc_title):
        request = f"ADD RFC {rfc_number} P2P-CI/1.0\r\nHost: {self.peer_ip}\r\nPort: {self.peer_port}\r\nTitle: {rfc_title}\r\n\r\n"
        self.server_connection.sendall(request.encode('utf-8'))
        response = self.server_connection.recv(4096).decode('utf-8')
        return response
    
    def send_lookup_request(self, rfc_number):
        request = f"LOOKUP RFC {rfc_number} P2P-CI/1.0\r\nHost: {self.peer_ip}\r\nPort: {self.peer_port}\r\n\r\n"
        self.server_connection.sendall(request.encode('utf-8'))
        response = self.server_connection.recv(4096).decode('utf-8')
        return response
    
    def send_list_request(self):
        request = f"LIST ALL P2P-CI/1.0\r\nHost: {self.peer_ip}\r\nPort: {self.peer_port}\r\n\r\n"
        self.server_connection.sendall(request.encode('utf-8'))
        response = self.server_connection.recv(4096).decode('utf-8')
        return response

    def send_leave_request(self):
        request = f"LEAVE P2P-CI/1.0\r\nHost: {self.peer_ip}\r\nPort: {self.peer_port}\r\n\r\n"
        self.server_connection.sendall(request.encode('utf-8'))
        response = self.server_connection.recv(4096).decode('utf-8')
        return response

    def close_connection_to_server(self):
        self.server_connection.close()

    @staticmethod
    def _parse_server_response_lines(response):
        lines = response.split("\r\n")
        status_line = lines[0].strip()
        parts = status_line.split()
        try:
            status_code = int(parts[1])
            phrase = " ".join(parts[2:])
        except Exception:
            return 400, "Bad Response", []
        body_lines = [ln for ln in lines[2:] if ln.strip()]
        return status_code, phrase, body_lines

    @staticmethod
    def _parse_list_lookup_body(lines):
        results = []
        for ln in lines:
            parts = ln.split()
            if len(parts) < 5:
                continue
            try:
                num = parts[1]
                host = parts[-2]
                port = int(parts[-1])
                title = " ".join(parts[2:-2])
                results.append((num, title, host, port))
            except Exception:
                continue
        return results

    @staticmethod
    def _recv_all(sock, bufsize=4096):
        sock.settimeout(2.0)
        chunks = []
        while True:
            try:
                data = sock.recv(bufsize)
            except socket.timeout:
                break
            if not data:
                break
            chunks.append(data)
            if len(data) < bufsize:
                break
        return b"".join(chunks)

    def download_rfc(self, rfc_number, save_dir="downloads", re_register=True):
        lookup_resp = self.send_lookup_request(rfc_number)
        code, phrase, lines = self._parse_server_response_lines(lookup_resp)
        if code != 200:
            return None, None
        print(f'Response of LOOKUP request: \n{lookup_resp}')
        candidates = self._parse_list_lookup_body(lines)
        if not candidates:
            return None, None

        os.makedirs(save_dir, exist_ok=True)
        content = None
        title = None
        chosen = None
        for num, t, host, port in candidates:
            if str(num) != str(rfc_number):
                continue
            try:
                req = (
                    f"GET RFC {rfc_number} P2P-CI/1.0\r\n"
                    f"Host: {self.peer_ip}\r\n"
                    f"OS: {platform.system()} {platform.release()}\r\n\r\n"
                )
                s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                s.settimeout(3.0)
                s.connect((host, int(port)))
                s.sendall(req.encode("utf-8"))
                raw = self._recv_all(s)
                s.close()
                header_body_split = raw.split(b"\r\n\r\n", 1)
                if not header_body_split:
                    continue
                headers = header_body_split[0].decode("utf-8", errors="ignore")
                print(f'Received Response Headers: \n{headers}\n')
                body = header_body_split[1] if len(header_body_split) > 1 else b""
                status_line = headers.split("\r\n")[0]
                try:
                    status_code = int(status_line.split()[1])
                except Exception:
                    continue
                if status_code != 200:
                    continue
                content = body
                title = t
                chosen = (host, port)
                break
            except Exception:
                continue

        if content is None:
            return None, None

        out_path = os.path.join(save_dir, f"{rfc_number}.txt")
        with open(out_path, "wb") as f:
            f.write(content)

        if re_register and title:
            try:
                self.send_add_request(rfc_number, title)
            except Exception:
                pass

        return out_path, chosen
    
    
def _load_peer_details(peer_number):
    path = os.path.join('peer_details', f'peer{peer_number}.txt')
    if not os.path.isfile(path):
        raise FileNotFoundError(f"Peer details file not found: {path}")
    with open(path, 'r', encoding='utf-8') as f:
        lines = [ln.strip() for ln in f.readlines() if ln.strip()]
    if len(lines) < 2:
        raise ValueError(f"Invalid peer details in {path}; expected at least IP and port")
    ip = lines[0]
    port = int(lines[1])
    rfcs = []
    for ln in lines[2:]:
        parts = ln.split(' ', 1)
        if not parts:
            continue
        number = parts[0]
        title = parts[1] if len(parts) > 1 else ''
        rfcs.append((number, title))
    return ip, port, rfcs


def _find_peer_by_port(peer_port):
    base = 'peer_details'
    if not os.path.isdir(base):
        return None
    for name in os.listdir(base):
        if not name.startswith('peer') or not name.endswith('.txt'):
            continue
        path = os.path.join(base, name)
        try:
            with open(path, 'r', encoding='utf-8') as f:
                lines = [ln.strip() for ln in f.readlines() if ln.strip()]
            if len(lines) >= 2 and int(lines[1]) == int(peer_port):
                return lines[0], int(lines[1])
        except Exception:
            continue
    return None


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="P2S client: register/list and download RFCs")
    parser.add_argument("--server-ip", default="127.0.0.1")
    parser.add_argument("--server-port", type=int, default=7734)
    parser.add_argument("--peer", type=int, help="Peer number to load from peer_details/peer<N>.txt for identity")
    parser.add_argument("--peer-ip", help="Override this peer's IP for P2S headers")
    parser.add_argument("--peer-port", type=int, help="Override this peer's upload port for P2S headers")
    parser.add_argument("--download", type=str, help="RFC number to download via LOOKUP->GET")
    parser.add_argument("--save-dir", default="downloads")
    parser.add_argument("--no-reregister", action="store_true", help="Do not ADD after download")
    parser.add_argument("--bulk-register", action="store_true", help="Load peer_details/* and register all RFCs")
    parser.add_argument("--register", action="store_true", help="Register this peer's RFCs from peer_details/peer<N>.txt")
    parser.add_argument("--list-all", action="store_true", help="List all the RFCs registered with the server")
    args = parser.parse_args()

    try:
        if args.download:
            identity_ip = None
            identity_port = None
            if args.peer is not None:
                identity_ip, identity_port, _ = _load_peer_details(args.peer)
            if identity_ip is None or identity_port is None:
                if args.peer_ip and args.peer_port:
                    identity_ip, identity_port = args.peer_ip, args.peer_port
                elif args.peer_port and not args.peer_ip:
                    found = _find_peer_by_port(args.peer_port)
                    if found:
                        identity_ip, identity_port = found
                if identity_ip is None or identity_port is None:
                    identity_ip = args.peer_ip or '127.0.0.1'
                    identity_port = args.peer_port or 4439

            client = PeerToServer(identity_ip, identity_port, args.server_ip, args.server_port)
            client.get_connection_to_server()
            path, chosen = client.download_rfc(args.download, save_dir=args.save_dir, re_register=not args.no_reregister)
            if path and chosen:
                print(f"Downloaded RFC {args.download} to {path} from {chosen[0]}:{chosen[1]}")
            else:
                print(f"RFC {args.download} not found in index or download failed.")
            client.close_connection_to_server()
        elif args.register:
            if args.peer is None:
                raise ValueError("Peer number must be specified with --peer to register RFCs.")
            identity_ip, identity_port, rfcs = _load_peer_details(args.peer)
            client = PeerToServer(identity_ip, identity_port, args.server_ip, args.server_port)
            client.get_connection_to_server()
            for rfc_number, rfc_title in rfcs:
                response = client.send_add_request(rfc_number, rfc_title)
                print(f"ADD RFC {rfc_number} Response:\n{response}\n")
            client.close_connection_to_server()
        elif args.bulk_register:
            number_of_peers = 5
            server_hostname = args.server_ip
            server_port = args.server_port
            for i in range(1, number_of_peers + 1):
                file = open(f"peer_details/peer{i}.txt", "r").readlines()
                peer_hostname = file[0].strip()
                peer_port = int(file[1].strip())
                peer = PeerToServer(peer_hostname, peer_port, server_hostname, server_port)
                peer.get_connection_to_server()
                for line in file[2:]:
                    rfc_number, rfc_title = line.strip().split(' ', 1)
                    _ = peer.send_add_request(rfc_number, rfc_title)
                if i == number_of_peers:
                    response = peer.send_list_request()
                    print(f"LIST Response for all RFCs added:\n{response}\n")
                peer.close_connection_to_server()
        elif args.list_all:
            if args.peer is None:
                raise ValueError("Peer number must be specified with --peer to make the query to the server")
            identity_ip, identity_port, rfcs = _load_peer_details(args.peer)
            client = PeerToServer(identity_ip, identity_port, args.server_ip, args.server_port)
            client.get_connection_to_server()
            response = client.send_list_request()
            print(f'Response received for LIST ALL request: \n{response}')
        else:
            print("No action specified. Use --download <RFC> or --bulk-register")
    except Exception as e:
        print(f"An error occurred: {e}")
        traceback.print_exc()