diff --git a/cmp.py b/cmp.py index 0b3b062..b035670 100644 --- a/cmp.py +++ b/cmp.py @@ -4,6 +4,7 @@ import random import string import json +import traceback from opcodes import OPCODES from strings import STRINGS @@ -40,56 +41,86 @@ async def handle_connection(self, reader: asyncio.StreamReader, writer: asyncio. client_address = writer.get_extra_info('peername') while True: - data = await reader.read(1024) - if not data: - continue try: - parsed = json.loads(data.decode()) - opcode = parsed['opcode'] - except: - writer.write(json.dumps({'opcode': OPCODES['PARSE_ERROR']}).encode()) - continue + data = await reader.read(1024) + if not data: + continue + try: + parsed = json.loads(data.decode()) + except: + writer.write(json.dumps({'opcode': OPCODES['PARSE_ERROR']}).encode()) + continue - if opcode not in OPCODES.values(): - writer.write(json.dumps({'opcode': OPCODES['UNKNOWN_OPCODE']}).encode()) - continue + opcode = parsed.get('opcode') + if opcode is None: # 0 equals to False or None and it gives an error + writer.write(json.dumps({'opcode': OPCODES['PARSE_ERROR']}).encode()) + continue - if opcode == OPCODES['CLIENT_DISCONNECT']: - writer.write(json.dumps({'opcode': OPCODES['CONNECTION_CLOSED']}).encode()) - writer.close() - await writer.wait_closed() - return - - if opcode == OPCODES['CLIENT_INITIALIZE'] or opcode == OPCODES['PING']: - writer.write(json.dumps({'opcode': OPCODES['CONNECTION_INITIALIZED']}).encode()) - continue + if opcode not in OPCODES.values(): + writer.write(json.dumps({'opcode': OPCODES['UNKNOWN_OPCODE']}).encode()) + continue - if opcode == OPCODES['IS_AVAILABLE']: - available = await self.is_available(parsed['address']) - writer.write(json.dumps({'opcode': OPCODES['IS_AVAILABLE'], 'result': available}).encode()) - continue + if opcode == OPCODES['CLIENT_DISCONNECT']: + writer.write(json.dumps({'opcode': OPCODES['CONNECTION_CLOSED']}).encode()) + writer.close() + await writer.wait_closed() + return + + if opcode == OPCODES['CLIENT_INITIALIZE'] or opcode == OPCODES['PING']: + writer.write(json.dumps({'opcode': OPCODES['CONNECTION_INITIALIZED']}).encode()) + continue - if opcode == OPCODES['REGISTER']: - result = await self.register_address(parsed['address'], parsed['password']) - writer.write(json.dumps({'opcode': OPCODES['REGISTER'], 'result': result}).encode()) - continue + if opcode == OPCODES['IS_AVAILABLE']: + address = parsed.get('address') + if not address: + writer.write(json.dumps({'opcode': OPCODES['PARSE_ERROR']}).encode()) + continue + available = await self.is_available(address.lower()) + writer.write(json.dumps({'opcode': OPCODES['IS_AVAILABLE'], 'result': available}).encode()) + continue - if opcode == OPCODES['SEND_MAIL']: - result = await self.send_mail(parsed['address'], parsed['password'], parsed['to_address'], parsed['text'], parsed['files']) - writer.write(json.dumps({'opcode': OPCODES['SEND_MAIL'], 'result': result}).encode()) - continue - - if opcode == OPCODES['GET_MAILS']: - result = await self.get_mails(parsed['address'], parsed['password']) - if result: - writer.write(json.dumps({'opcode': OPCODES['GET_MAILS'], 'result': True, 'data': result}).encode()) + if opcode == OPCODES['REGISTER']: + address = parsed.get('address') + password = parsed.get('password') + if not address or not password: + writer.write(json.dumps({'opcode': OPCODES['PARSE_ERROR']}).encode()) + continue + result = await self.register_address(address.lower(), password) + writer.write(json.dumps({'opcode': OPCODES['REGISTER'], 'result': result}).encode()) continue - writer.write(json.dumps({'opcode': OPCODES['GET_MAILS'], 'result': False}).encode()) - continue - if opcode == OPCODES['UPLOAD_FILE']: - ... - #ill do it tomorrow im tired + if opcode == OPCODES['SEND_MAIL']: + address = parsed.get('address') + password = parsed.get('password') + to_address = parsed.get('to_address') + text = parsed.get('text') + files = parsed.get('files') + if not address or not password or not to_address or not text or not files: + writer.write(json.dumps({'opcode': OPCODES['PARSE_ERROR']}).encode()) + continue + result = await self.send_mail(address.lower(), password, to_address.lower(), text, files) + writer.write(json.dumps({'opcode': OPCODES['SEND_MAIL'], 'result': result}).encode()) + continue + + if opcode == OPCODES['GET_MAILS']: + address = parsed.get('address') + password = parsed.get('password') + if not address or not password: + writer.write(json.dumps({'opcode': OPCODES['PARSE_ERROR']}).encode()) + continue + result = await self.get_mails(address.lower(), password) + if result: + writer.write(json.dumps({'opcode': OPCODES['GET_MAILS'], 'result': True, 'data': result}).encode()) + continue + writer.write(json.dumps({'opcode': OPCODES['GET_MAILS'], 'result': False}).encode()) + continue + + if opcode == OPCODES['UPLOAD_FILE']: + ... + #ill do it tomorrow im tired + except: + traceback.print_exc() + writer.write(json.dumps({'opcode': OPCODES['PARSE_ERROR']}).encode()) async def is_available(self, address: str) -> bool: """ @@ -98,18 +129,18 @@ async def is_available(self, address: str) -> bool: - address: Address (str) Returns: bool """ - if address.lower() in self.addresses.keys(): - return {'result': False, 'message': STRINGS['ADDRESS_NOT_AVAILABLE']} + if address in self.addresses.keys(): + return {'opcode': OPCODES['IS_AVAILABLE'], 'result': False, 'message': STRINGS['ADDRESS_NOT_AVAILABLE']} if len(address) > self.maximum_address_length: - return {'result': False, 'message': STRINGS['ADDRESS_TOO_LONG']} + return {'opcode': OPCODES['IS_AVAILABLE'], 'result': False, 'message': STRINGS['ADDRESS_TOO_LONG']} if len(address) < 3: - return {'result': False, 'message': STRINGS['ADDRESS_TOO_SHORT']} + return {'opcode': OPCODES['IS_AVAILABLE'], 'result': False, 'message': STRINGS['ADDRESS_TOO_SHORT']} for k in address: if k not in self.allowed_address_characters: return {'result': False, 'message': STRINGS['ADDRESS_IS_BAD']} - return {'result': True, 'message': STRINGS['ADDRESS_AVAILABLE']} + return {'opcode': OPCODES['IS_AVAILABLE'], 'result': True, 'message': STRINGS['ADDRESS_AVAILABLE']} async def register_address(self, address: str, password: str) -> bool: """ @@ -121,15 +152,16 @@ async def register_address(self, address: str, password: str) -> bool: """ available = await self.is_available(address) if not available['result']: + available['opcode'] = OPCODES['REGISTER'] return available - self.addresses[address.lower()] = { + self.addresses[address] = { 'password': password, 'mails': [], 'register_date': time.time(), 'admin': False } - return {'result': True, 'message': STRINGS['REGISTER_SUCCESSFUL']} + return {'opcode': OPCODES['REGISTER'], 'result': True, 'message': STRINGS['REGISTER_SUCCESSFUL']} async def check_credentials(self, address: str, password: str) -> bool: """ @@ -156,23 +188,23 @@ async def send_mail(self, address: str, password: str, to_address: str, text: st - files: Files to send (list[dict], for example: [{"file_id": some_id}, ...] or [] if no files) Returns: bool """ - is_valid = await self.check_credentials(address.lower(), password) + is_valid = await self.check_credentials(address, password) if not is_valid: - return {'result': False, 'message': STRINGS['INVALID_CREDENTIALS']} + return {'opcode': OPCODES['SEND_MAIL'], 'result': False, 'message': STRINGS['INVALID_CREDENTIALS']} if to_address not in self.addresses.keys(): - return {'result': False, 'message': STRINGS['ADDRESS_NOT_FOUND']} + return {'opcode': OPCODES['SEND_MAIL'], 'result': False, 'message': STRINGS['ADDRESS_NOT_FOUND']} if len(files) > 15: - return {'result': False, 'message': STRINGS['FILES_LIMIT']} + return {'opcode': OPCODES['SEND_MAIL'], 'result': False, 'message': STRINGS['FILES_LIMIT']} for file in files: if not file: continue if 'file_id' not in file.keys() or len(file.keys()) > 1: - return {'result': False, 'message': STRINGS['INVALID_FILE']} + return {'opcode': OPCODES['SEND_MAIL'], 'result': False, 'message': STRINGS['INVALID_FILE']} - self.addresses[to_address.lower()]['mails'].append({ + self.addresses[to_address]['mails'].append({ 'out': False, 'to_address': None, 'from_address': address, @@ -180,7 +212,7 @@ async def send_mail(self, address: str, password: str, to_address: str, text: st 'files': files, 'sent_at': time.time() }) - self.addresses[address.lower()]['mails'].append({ + self.addresses[address]['mails'].append({ 'out': True, 'to_address': to_address, 'from_address': None, @@ -188,7 +220,7 @@ async def send_mail(self, address: str, password: str, to_address: str, text: st 'files': files, 'sent_at': time.time() }) - return {'result': True, 'message': STRINGS['MAIL_SENT']} + return {'opcode': OPCODES['SEND_MAIL'], 'result': True, 'message': STRINGS['MAIL_SENT']} async def get_mails(self, address: str, password: str) -> list[dict] | bool: """ @@ -198,11 +230,11 @@ async def get_mails(self, address: str, password: str) -> list[dict] | bool: - password: Password (str, for example: verysecurepassword) Returns: list[dict] | bool (on error) """ - is_valid = await self.check_credentials(address.lower(), password) + is_valid = await self.check_credentials(address, password) if not is_valid: - return {'result': False, 'message': STRINGS['INVALID_CREDENTIALS']} + return {'opcode': OPCODES['GET_MAILS'], 'result': False, 'message': STRINGS['INVALID_CREDENTIALS']} - return json.dumps(self.addresses[address.lower()]['mails']) + return {'opcode': OPCODES['GET_MAILS'], 'result': True, 'data': self.addresses[address]['mails']} class Client: def __init__(self, host: str, port: int) -> None: @@ -361,4 +393,4 @@ async def get_mails(self, address: str, password: str, timeout: float = 5) -> li mails = parsed['data'] except: return False - return mails \ No newline at end of file + return mails diff --git a/strings.py b/strings.py index 36c6b1b..b50d393 100644 --- a/strings.py +++ b/strings.py @@ -6,12 +6,12 @@ 'ADDRESS_AVAILABLE': 'The address is available.', 'ADDRESS_NOT_FOUND': 'The address you provided was not found.', - 'REGISTER_SUCCESSFUL': 'The address was successfuly registered.', + 'REGISTER_SUCCESSFUL': 'The address was successfully registered.', - 'INVALID_CREDENTIALS': 'The credentials you provided are wrong.', + 'INVALID_CREDENTIALS': 'The credentials you provided are incorrect.', - 'FILES_LIMIT': 'You reached the files limit.', + 'FILES_LIMIT': 'You have reached the files limit.', 'INVALID_FILE': 'You provided invalid file data.', - 'MAIL_SENT': 'The mail was successfuly sent.', -} \ No newline at end of file + 'MAIL_SENT': 'The mail was successfully sent.', +} diff --git a/test.py b/test.py index 2c91272..c63d24b 100644 --- a/test.py +++ b/test.py @@ -2,40 +2,61 @@ import asyncio import base64 from opcodes import OPCODES +import json +import sys async def main(): server = cmp.Server('127.0.0.1', 16760) client = cmp.Client('127.0.0.1', 16760) - await server.start() - await client.connect() - - available = await client.is_address_available('admin', 1) # True - print(available) - - result = await client.register_address('admin', 'admin', 1) # True - print(result) - - result = await client.register_address('admin2', 'admin2', 1) # True - print(result) - - with open('image.png', 'rb') as f: - image_content = f.read() - - result = await client.send_mail('admin', 'admin', 'admin2', 'Hey!', files=[{'file_id': 'image.png'}]) # True - print(result) - - - #mails = await client.get_mails('admin', 'admin') # - #print(mails) - - #await client.send_raw_message(OPCODES['PING']) - #result = await client.wait_for_raw_message(1024, 1) - #print(f'Message: {result}') - - - - - await client.close() + try: + await server.start() + await asyncio.sleep(0.05) + await client.connect() + + available = await client.is_address_available('admin', 1) + print('is "admin" available:', available) + + result = await client.register_address('admin', 'admin', 1) + print('register admin:', result) + + result = await client.register_address('admin2', 'admin2', 1) + print('register admin2:', result) + + with open('image.png', 'rb') as f: + image_content = f.read() + + result = await client.send_mail('admin', 'admin', 'admin2', 'Hey!', files=[{'file_id': 'image.png'}]) + print('send mail result:', result) + + mails_admin2 = await client.get_mails('admin2', 'admin2', 1) + print('admin2 mails raw:', mails_admin2) + try: + print('admin2 mails parsed:', json.dumps(mails_admin2, indent=2)) + except Exception: + pass + + mails_admin = await client.get_mails('admin', 'admin', 1) + print('admin mails raw:', mails_admin) + try: + print('admin mails parsed:', json.dumps(mails_admin, indent=2)) + except Exception: + pass + + await client.close() + except Exception as e: + print('Error during test run:', e, file=sys.stderr) + try: + await client.close() + except Exception: + pass + finally: + try: + srv = getattr(server, 'server', None) + if srv: + srv.close() + await srv.wait_closed() + except Exception: + pass if __name__ == '__main__': - asyncio.run(main()) \ No newline at end of file + asyncio.run(main())