diff --git a/.env.example b/.env.example index 3a8a89d..ec0415c 100644 --- a/.env.example +++ b/.env.example @@ -10,4 +10,6 @@ ACCESS_TOKEN="xxxxxxxxxxxxxxxxxxxxx" # Optional, use user_id and password is rec BARD_TOKEN="xxxxxxxxxxxxxxxxxxxx", # Optional, for !bard command JAILBREAKENABLED="true" # Optional BING_AUTH_COOKIE="xxxxxxxxxxxxxxxxxxx" # _U cookie, Optional, for Bing Image Creator -MARKDOWN_FORMATTED="true" # Optional \ No newline at end of file +MARKDOWN_FORMATTED="true" # Optional +IMPORT_KEYS_PATH="element-keys.txt" # Optional +IMPORT_KEYS_PASSWORD="xxxxxxx" # Optional \ No newline at end of file diff --git a/bot.py b/bot.py index ade0edf..176097b 100644 --- a/bot.py +++ b/bot.py @@ -3,11 +3,11 @@ import os import re import sys import traceback -from typing import Union +from typing import Union, Optional import aiohttp from nio import (AsyncClient, AsyncClientConfig, InviteMemberEvent, JoinError, - KeyVerificationCancel, KeyVerificationEvent, + KeyVerificationCancel, KeyVerificationEvent, EncryptionError, KeyVerificationKey, KeyVerificationMac, KeyVerificationStart, LocalProtocolError, LoginResponse, MatrixRoom, MegolmEvent, RoomMessageText, ToDeviceError) @@ -42,6 +42,8 @@ class Bot: jailbreakEnabled: Union[bool, None] = True, bing_auth_cookie: Union[str, None] = '', markdown_formatted: Union[bool, None] = False, + import_keys_path: Optional[str] = None, + import_keys_password: Optional[str] = None, ): if (homeserver is None or user_id is None or device_id is None): @@ -61,6 +63,8 @@ class Bot: self.room_id = room_id self.api_key = api_key self.chatgpt_api_endpoint = chatgpt_api_endpoint + self.import_keys_path = import_keys_path + self.import_keys_password = import_keys_password self.session = aiohttp.ClientSession() @@ -173,7 +177,7 @@ class Bot: raw_user_message = event.body # print info to console - print( + logger.info( f"Message received in room {room.display_name}\n" f"{room.user_name(event.sender)} | {raw_user_message}" ) @@ -352,14 +356,12 @@ class Bot: if "emoji" not in event.short_authentication_string: estr = ("Other device does not support emoji verification " f"{event.short_authentication_string}. Aborting.") - print(estr) logger.info(estr) return resp = await client.accept_key_verification( event.transaction_id) if isinstance(resp, ToDeviceError): estr = f"accept_key_verification() failed with {resp}" - print(estr) logger.info(estr) sas = client.key_verifications[event.transaction_id] @@ -368,7 +370,6 @@ class Bot: resp = await client.to_device(todevice_msg) if isinstance(resp, ToDeviceError): estr = f"to_device() failed with {resp}" - print(estr) logger.info(estr) elif isinstance(event, KeyVerificationCancel): # anytime @@ -391,7 +392,6 @@ class Bot: # We only need to inform the user. estr = (f"Verification has been cancelled by {event.sender} " f"for reason \"{event.reason}\".") - print(estr) logger.info(estr) elif isinstance(event, KeyVerificationKey): # second step @@ -409,7 +409,7 @@ class Bot: """ sas = client.key_verifications[event.transaction_id] - print(f"{sas.get_emoji()}") + logger.info(f"{sas.get_emoji()}") # don't log the emojis # The bot process must run in forground with a screen and @@ -422,36 +422,30 @@ class Bot: if yn.lower() == "y": estr = ("Match! The verification for this " "device will be accepted.") - print(estr) logger.info(estr) resp = await client.confirm_short_auth_string( event.transaction_id) if isinstance(resp, ToDeviceError): estr = ("confirm_short_auth_string() " f"failed with {resp}") - print(estr) logger.info(estr) elif yn.lower() == "n": # no, don't match, reject estr = ("No match! Device will NOT be verified " "by rejecting verification.") - print(estr) logger.info(estr) resp = await client.cancel_key_verification( event.transaction_id, reject=True) if isinstance(resp, ToDeviceError): estr = (f"cancel_key_verification failed with {resp}") - print(estr) logger.info(estr) else: # C or anything for cancel estr = ("Cancelled by user! Verification will be " "cancelled.") - print(estr) logger.info(estr) resp = await client.cancel_key_verification( event.transaction_id, reject=False) if isinstance(resp, ToDeviceError): estr = (f"cancel_key_verification failed with {resp}") - print(estr) logger.info(estr) elif isinstance(event, KeyVerificationMac): # third step @@ -478,13 +472,11 @@ class Bot: estr = (f"Cancelled or protocol error: Reason: {e}.\n" f"Verification with {event.sender} not concluded. " "Try again?") - print(estr) logger.info(estr) else: resp = await client.to_device(todevice_msg) if isinstance(resp, ToDeviceError): estr = f"to_device failed with {resp}" - print(estr) logger.info(estr) estr = (f"sas.we_started_it = {sas.we_started_it}\n" f"sas.sas_accepted = {sas.sas_accepted}\n" @@ -492,7 +484,6 @@ class Bot: f"sas.timed_out = {sas.timed_out}\n" f"sas.verified = {sas.verified}\n" f"sas.verified_devices = {sas.verified_devices}\n") - print(estr) logger.info(estr) estr = ("Emoji verification was successful!\n" "Initiate another Emoji verification from " @@ -500,16 +491,13 @@ class Bot: "Or if done verifying, hit Control-C to stop the " "bot in order to restart it as a service or to " "run it in the background.") - print(estr) logger.info(estr) else: estr = (f"Received unexpected event type {type(event)}. " f"Event is {event}. Event will be ignored.") - print(estr) logger.info(estr) except BaseException: estr = traceback.format_exc() - print(estr) logger.info(estr) # !chat command @@ -633,11 +621,22 @@ class Bot: resp = await self.client.login(password=self.password) if not isinstance(resp, LoginResponse): logger.error("Login Failed") - print(f"Login Failed: {resp}") sys.exit(1) except Exception as e: logger.error(f"Error: {e}", exc_info=True) + # import keys + async def import_keys(self): + resp = await self.client.import_keys( + self.import_keys_path, + self.import_keys_password + ) + if isinstance(resp, EncryptionError): + logger.error(f"import_keys failed with {resp}") + else: + logger.info( + f"import_keys success, please remove import_keys configuration!!!") + # sync messages in the room async def sync_forever(self, timeout=30000, full_state=True) -> None: diff --git a/config.json.sample b/config.json.sample index af351a8..d50dc85 100644 --- a/config.json.sample +++ b/config.json.sample @@ -10,5 +10,7 @@ "access_token": "xxxxxxx", "bard_token": "xxxxxxx", "bing_auth_cookie": "xxxxxxxxxxx", - "markdown_formatted": true + "markdown_formatted": true, + "import_keys_path": "element-keys.txt", + "import_keys_password": "xxxxxxxxx" } diff --git a/log.py b/log.py index 6beb356..249efb4 100644 --- a/log.py +++ b/log.py @@ -16,10 +16,11 @@ def getlogger(): # create formatters warn_format = logging.Formatter( - '%(name)s - %(funcName)s - %(levelname)s - %(message)s') + '%(asctime)s - %(funcName)s - %(levelname)s - %(message)s') error_format = logging.Formatter( '%(asctime)s - %(name)s - %(funcName)s - %(levelname)s - %(message)s') - info_format = logging.Formatter('%(message)s') + info_format = logging.Formatter( + '%(asctime)s - %(levelname)s - %(message)s') # set formatter warn_handler.setFormatter(warn_format) diff --git a/main.py b/main.py index b71e5d8..8338bf1 100644 --- a/main.py +++ b/main.py @@ -8,7 +8,7 @@ logger = getlogger() async def main(): - + need_import_keys = False if os.path.exists('config.json'): fp = open('config.json', 'r', encoding="utf8") config = json.load(fp) @@ -25,7 +25,12 @@ async def main(): jailbreakEnabled=config.get('jailbreakEnabled'), bing_auth_cookie=config.get('bing_auth_cookie'), markdown_formatted=config.get('markdown_formatted'), + import_keys_path=config.get('import_keys_path'), + import_keys_password=config.get( + 'import_keys_password'), ) + if config.get('import_keys_path') and config.get('import_keys_password') is not None: + need_import_keys = True else: matrix_bot = Bot(homeserver=os.environ.get('HOMESERVER'), @@ -42,10 +47,17 @@ async def main(): bing_auth_cookie=os.environ.get("BING_AUTH_COOKIE"), markdown_formatted=os.environ.get( "MARKDOWN_FORMATTED", "false").lower() in ('true', '1', 't'), + import_keys_path=os.environ.get("IMPORT_KEYS_PATH"), + import_keys_password=os.environ.get( + "IMPORT_KEYS_PASSWORD"), ) + if os.environ.get("IMPORT_KEYS_PATH") and os.environ.get("IMPORT_KEYS_PASSWORD") is not None: + need_import_keys = True await matrix_bot.login() - + if need_import_keys: + logger.info("start import_keys process, this may take a while...") + await matrix_bot.import_keys() await matrix_bot.sync_forever(timeout=30000, full_state=True)