Socket Programming (Domain Name Service)

CODE-K·2025년 2월 14일
0

컴퓨터 네트워크

목록 보기
16/16

Database를 통해 Domain과 IP주소를 관리하고 DB를 활용하여 Client측의 요청을 처리하는 프로그램입니다.

Database는 Python의 sqlite3를 사용하였습니다.

dns.py

import sqlite3, json
from typing import TypedDict, Optional
from enum import Enum

class DNS():
    def __init__(self) -> None:
        self.con = sqlite3.connect('dns.db')
        self.cur = self.con.cursor()

        self.cur.execute("""
        CREATE TABLE IF NOT EXISTS domains (ip text PRIMARY KEY, dname text UNIQUE)
        """)

    def insert_domain(self, ip: str, dname: str):
        self.cur.execute(f"INSERT INTO domains VALUES (?, ?)", (ip, dname))
        self.con.commit()

    def delete_domain(self, ip: str, dname: str):
        self.cur.execute(f"DELETE FROM domains WHERE ip = ? AND dname = ?", (ip, dname))
        self.con.commit()

    def search_ip(self, ip: str):
        '''Search using ip'''
        res = self.cur.execute(f"SELECT dname FROM domains WHERE ip=?", (ip,))  # tuple
        data = res.fetchone()
        if data:
            return data[0]
        else:
            return None
    
    def search_dname(self, domain: str):
        '''Search using dname'''
        res = self.cur.execute(f"SELECT ip FROM domains WHERE dname=?", (domain,))  # tuple
        data = res.fetchone()
        if data:
            return data[0]
        else:
            return None
    
    def close(self):
        self.con.close()

class DataField(TypedDict, total=False):
    type: str  # A, NS, CNAME etc
    ip: str
    dname: str

class MyDNSProtocol(TypedDict):
    '''
    status flag
    Client
    1: insert
    2: delete (matching ip, dname)
    3: search ip
    4: search dname
    

    Server
    1: ok
    
    11: Failed Insert (duplicate data exists)
    12: Failed Delete (non-matching ip,dname)
    13: Failed Search (ip/dname was not given)
    14: Failed Search (not found)

    20: Unknown Request
    '''
    status: int
    data: DataField

def read_data(data):
    json_data = json.loads(data.decode('utf-8'))  # byte string -> dict
    return MyDNSProtocol(json_data)

def parse_data(status: int, type: str='A', ip: str=None, dname: str=None):
    data_field = DataField(type=type, ip=ip, dname=dname)
    payload = MyDNSProtocol(status=status, data=data_field)
    return json.dumps(payload, ensure_ascii=False).encode('utf-8')

def check_db(dns: DNS):
    for row in dns.cur.execute("""SELECT * FROM domains"""):
        print(row)
     

if __name__ == "__main__":
    dns = DNS()
    try:
        pass
        # dns.insert_domain('1.2.3.4', 'test.domain')
        # dns.insert_domain('1.2.3.5', 'test.domain.2')
        # dns.insert_domain('1.2.3.6', 'test.domain.3')
        # dns.insert_domain('1.2.3.7', 'test.domain')
        # domain = dns.search_ip('1.2.3.4')
        # ip = dns.search_dname('test.domain')
    except sqlite3.IntegrityError as e:
        print(f"Error: {e}")

    check_db(dns)
    input()

Client와 Server에서 모두 사용하는 dns.py 코드입니다.

Database 사용을 위해 Python 자체 sqlite3를 사용하였고, 데이터 저장을 위해 json,
TypedDict 등을 사용하였습니다.


sqlite3.connect를 통해 dns.db와 연동하였고, cursor를 생성합니다.

execute를 통해 Database의 테이블을 생성
ip, dname은 이미 존재하는 내용의 중복 방지를 위해 PRIMARY KEY와 UNIQUE로 지정합니다.


기능 구현 부분입니다.

insert_domain : ip, dname을 받아 INSERT를 사용하여 각 값을 이용하여 삽입하고

delete_domain : ip와 dname을 받아 두 값이 모두 일치하는 경우 테이블에서 delete 합니다.

search_ip : ip를 입력 받아 해당되는 dname을 반환하는 것으로 일치하는 내용이 없을 경우 None을 반환합니다.

search_dname : dname을 입력받아 일치하는 ip를 반환하는 것으로 일치하지 않을 경우 None을 반환합니다.

close : DB cursor를 종료합니다.


DataField : TypedDict 형태를 이용하였고, Domain의 Type, ip, dname으로 되어있습니다.

MyDNSProtocol : 주석으로 status flag를 서술하였습니다.
Client측의 입력값에 대한 내용과 Server에서 반환하는 상태 정보에 대한 내용입니다.

status : int, data는 DataField 형태로 구성하였습니다.

read_data와 parse_data : json 형태의 데이터를 읽고 전송하기 위해 구현하였습니다.

구현 해놓은 DataField와 MyDNSProtocol을 이용하여 처리합니다.

check_db : DB 전체 내용을 확인하기 위해 출력합니다.

이후 main 부분은 dns에서 오류 확인을 위해 구현한 부분입니다.


client_dns.py

import socket
from dns import *

host = '0.0.0.0'
port = 12000

if __name__ == "__main__":
    host_input = input('Enter server IP (Leave empty for localhost): ')  # for non-local IP connections
    if host_input:
        host = host_input

    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.connect((host, port))
        print('Connected to server')

        while True:
            print('\nSelect Mode(Exit: 0, Insert: 1, Delete: 2, Search IP: 3, Search Domain: 4)')
            try:
                mode = int(input('Mode: '))
            except ValueError:
                continue

            if mode == 0:
                break
            ip = None
            dname = None

            # input ip
            if mode != 4:  # no need in search dname 
                while not ip:
                    ip = input('IP: ')

            # input dname
            if mode != 3:
                while not dname:
                    dname = input('Domain Name: ')

            request = parse_data(mode, ip=ip, dname=dname)
            s.sendall(request)
            raw_data = s.recv(2048)

            data = read_data(raw_data)

            if (status := data['status']) == 1:
                print('Success')
                if mode == 3 or mode == 4:
                    print(f"IP: {data['data']['ip']}, Domain Name: {data['data']['dname']}")

            elif status == 11:
                print('Failed insert due to duplicate entry')
            elif status == 12:
                print('Failed delete due to incorrect ip or dname')
            elif status == 13:
                print('Failed Insert due to empty search value')
            elif status == 14:  # entry not found
                print('Searched entry was not found')
            else:
                print('Unknown Request')

        s.close()
    input()

host_input : 연결한 Server의 주소를 input하고 input한 Server의 주소를 통해 생성한 Client socket을 연결합니다.

본격적인 내용은 while문 내부에 있습니다.

사용자에게 Mode 선택 문장을 출력하고 input을 통해 원하는 Mode를 선택합니다.

if문을 통해 mode를 다음과 같이 구분하고,

0 : Exit
1 : Insert
2 : Delete
3 : Search IP
4 : Search Domain

해당 mode에 대한 입력을 나누어 받습니다.

parse_data : dns.py에서 지정한 json 형식으로 반환합니다.

sendall : request를 server에 보내고, 요청에 대한 응답을 raw_data로 recv합니다.

응답을 data에 저장하고 data의 정보를 토대로 응답을 출력하게 됩니다.

status를 통해 응답을 나누는데, 이는 Server에서 보낸 데이터에 포함된 내용이며,

1 : 정상

10 이후의 status는 각각의 오류들에 해당합니다.

Client에서 0을 입력하면 Exit하여 Client socket을 close하고 종료합니다.


server_dns.py

import socket
from dns import *

host = '0.0.0.0'
port = 12000

def check_ip():
    '''check if IP is valid format'''
    pass

def check_domain():
    '''check if domain name is valid format'''
    pass

if __name__ == "__main__":
    dns_server = DNS()
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind((host, port))
        s.listen()

        print(f'Server listening on port {port}')
        while True:
            conn, addr = s.accept()
            print(f'Connected by {addr}')
            with conn:  # context manager to ensure exit
                while True:
                    # recieve data           
                    raw_data = conn.recv(2048)

                    # connection closed
                    if not raw_data:
                        print(f'Connection closed by {addr}')
                        break

                    # decode read data format
                    data = read_data(raw_data)
                    type = data['data']['type']  # unused
                    ip = data['data']['ip']
                    dname = data['data']['dname']

                    # input validation for potential SQL injections
                    # check_ip()
                    # check_domain()

                    # insert
                    if ((status := data.get('status')) == 1):
                        print('Insert Request: ', end='')
                        try:
                            dns_server.insert_domain(ip, dname)
                            print(f'Success - ({ip}: {dname})')
                            conn.sendall(parse_data(1))
                        except sqlite3.IntegrityError:
                            conn.sendall(parse_data(11))  # status: 11 = fail
                            print('Failed')
                    # delete
                    elif (status == 2):
                        print('Delete Request: ', end='')
                        if (dns_server.search_ip(ip) == dname):
                            dns_server.delete_domain(ip, dname)
                            print(f'Success - ({ip}: {dname})')
                            conn.sendall(parse_data(1))
                        else:
                            print('Fail')
                            conn.sendall(parse_data(12))
                    # search
                    elif (status == 3):
                        print('Search IP Request: ', end='')
                        if not ip:
                            print('Fail')
                            conn.sendall(parse_data(13))
                        else:
                            if (found_dname := dns_server.search_ip(ip)):
                                print(f'Success - ({ip}: {found_dname})')
                                conn.sendall(parse_data(1, ip=ip, dname=found_dname))
                            else: # found_dname == None
                                print('Failed (Not found)')
                                conn.sendall(parse_data(14))
                    elif (status == 4):
                        print('Search Domain Request: ', end='')
                        if not dname:
                            print('Fail')
                            conn.sendall(parse_data(13))
                        else:
                            if (found_ip := dns_server.search_dname(dname)):
                                print(f'Success - ({found_ip}: {dname})')
                                conn.sendall(parse_data(1, ip=found_ip, dname = dname))
                            else: # found_dname == None
                                print('Failed (Not found)')
                                conn.sendall(parse_data(14))
                    
                    else:
                        print('Unknown Request')
                        conn.sendall(parse_data(20))

Server측은 DNS를 활용하고 socket을 생성 후 bind, listen 합니다.

주요 기능은 while 내부의 내용입니다.

accept : 통해 Client의 connect를 연결

raw_data : Client 측에서 보낸 request (요청)을 recv합니다.

raw_data가 없을 경우 연결 실패로 Connection closed를 출력합니다.

raw_data, 즉 request가 존재할 경우 해당 내용을 읽어와 정보들을 분류합니다.

TypedDict를 사용하였기 때문에 type, ip, dname을 따로 읽어옵니다.


Client의 요청은 5가지 분류.

앞서 설명한 것과 같이

0 : Exit , 1 : Insert , 2 : Delete , 3 : Search ip , 4 : Search dname 입니다.

insert : insert_domain을 통해 DB에 insert하며,
성공 여부를 출력 parse_data를 통해 status 정보를 전송합니다.

만약 오류가 발생했을 경우, 해당하는 오류 내용을 전송합니다. (11번)

delete : ip와 dname 모두 일치하여야 삭제가 진행되기 때문에,
해당 내용을 점검하고 문제가 없을 경우 delete_domain을 진행, parse_data를 통해 status 정보를 전송합니다.

실패 시에는 해당하는 오류 정보를 전송합니다. (12번)

search_ip : ip 입력이 없다면 오류 정보를 전송합니다. (13번)

존재한다면 search_ip를 사용하여 대응하는 dname을 찾아 출력합니다.

만약 대응하는 dname이 없다면 오류 정보를 전송합니다. (14번)

search_dname : dname의 입력이 없다면 13번 오류 전송,
입력이 있다면 search_dname을 통해 대응하는 ip 검색 및 출력,
만약 대응하는 ip가 없다면 오류 정보를 전송합니다. (14번)

또한 초기 request가 0~4번이 아닌 다른 값이 입력되었다면 Unknown Request, 20번 오류 전송합니다.

profile
개발자 지망생입니다.

0개의 댓글