120 lines
4.1 KiB
Python
120 lines
4.1 KiB
Python
import shlex
|
|
import socket
|
|
import struct
|
|
import argparse
|
|
|
|
|
|
def print_logo():
|
|
print(r"""
|
|
╔════════════════════════════════════════════════════════════╗
|
|
║ _ _____ ______ _____ _____ _____ ║
|
|
║ | | | __ \| ____| __ \_ _|/ ____|║
|
|
║ _ __ ___ | |_ __ _ | |__) | |__ | | | || | | (___ ║
|
|
║| '_ \ / _ \| __| / _` | | _ /| __| | | | || | \___ \ ║
|
|
║| | | | (_) | |_ | (_| | | | \ \| |____| |__| || |_ ____) |║
|
|
║|_| |_|\___/ \__| \__,_| |_| \_\______|_____/_____|_____/ ║
|
|
╚════════════════════════════════════════════════════════════╝
|
|
""")
|
|
|
|
def recv_exact(sock: socket.socket, n: int) -> bytes:
|
|
data = b''
|
|
while len(data) < n:
|
|
chunk = sock.recv(n - len(data))
|
|
if len(chunk) == 0:
|
|
raise ConnectionError('EOF from server')
|
|
data += chunk
|
|
return data
|
|
|
|
def send_command(sock: socket.socket, text: str) -> None:
|
|
parts = shlex.split(text)
|
|
if not parts:
|
|
return
|
|
|
|
chunks = []
|
|
|
|
n_args = len(parts)
|
|
chunks.append(struct.pack('!I', n_args))
|
|
|
|
for part in parts:
|
|
part_bytes = part.encode()
|
|
chunks.append(struct.pack('!I', len(part_bytes)))
|
|
chunks.append(part_bytes)
|
|
|
|
payload = b''.join(chunks)
|
|
|
|
sock.sendall(struct.pack('!I', len(payload)) + payload)
|
|
|
|
def recv_response(sock: socket.socket) -> str:
|
|
header = recv_exact(sock, 4)
|
|
(length,) = struct.unpack('!I', header)
|
|
if length > 10_000_000:
|
|
raise ValueError('Message length is too long')
|
|
|
|
body = recv_exact(sock, length)
|
|
if length < 4:
|
|
raise ValueError("Response too short")
|
|
|
|
(status,) = struct.unpack('!I', body[:4])
|
|
msg = body[4:]
|
|
|
|
if status == 0:
|
|
if not msg:
|
|
return "(ok)"
|
|
return msg.decode("utf-8", errors='replace')
|
|
elif status == 2:
|
|
return "(nil)"
|
|
else:
|
|
return f"(err) {msg.decode('utf-8', errors='replace')}"
|
|
|
|
def main():
|
|
print_logo()
|
|
parser = argparse.ArgumentParser(description='NOT(Redis) client')
|
|
parser.add_argument('-H', '--host', type=str, required=False, default='127.0.0.1', help='Server host')
|
|
parser.add_argument('-P', '--port', type=int, required=False, default=6379, help='Server port')
|
|
parser.add_argument('-M', '--message', type=str, required=False, default=None, help='Command to send (e.g. "set k v")')
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
with socket.create_connection((args.host, args.port)) as sock:
|
|
if args.message:
|
|
send_command(sock, args.message)
|
|
response = recv_response(sock)
|
|
print(response)
|
|
else:
|
|
print(f"Connected to {args.host}:{args.port}")
|
|
print("Type 'quit' or 'exit' to leave.")
|
|
import readline
|
|
|
|
while True:
|
|
try:
|
|
prompt = f"{args.host}:{args.port}> "
|
|
cmd_line = input(prompt).strip()
|
|
|
|
if not cmd_line:
|
|
continue
|
|
|
|
if cmd_line.lower() in ('quit', 'exit'):
|
|
break
|
|
|
|
send_command(sock, cmd_line)
|
|
response = recv_response(sock)
|
|
print(response)
|
|
|
|
except KeyboardInterrupt:
|
|
print("\nType 'quit' or 'exit' to leave.")
|
|
except EOFError:
|
|
break
|
|
|
|
except ConnectionRefusedError as e:
|
|
print('Connection refused by server:', e)
|
|
exit(1)
|
|
except ConnectionError as e:
|
|
print('Connection error:', e)
|
|
exit(1)
|
|
except Exception as e:
|
|
print('Unexpected error:', e)
|
|
exit(1)
|
|
|
|
if __name__ == '__main__':
|
|
main()
|