bobmaster/matrix_chatgpt_bot
1
0
Fork 0
Code Issues 5 Pull requests 2 Projects Releases 10 Packages Wiki Activity

feat: import E2E room keys

Browse source
This commit is contained in:
hibobmaster 2023-04-20 15:39:14 +08:00
parent cb0c913547
commit 7fecffad54
Signed by: bobmaster
GPG key ID: 316B77D7914D713C
5 changed files with 43 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
.env.example
View file

@ -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
MARKDOWN_FORMATTED="true" # Optional
IMPORT_KEYS_PATH="element-keys.txt" # Optional
IMPORT_KEYS_PASSWORD="xxxxxxx" # Optional

41
bot.py
View file

@ -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:

4
config.json.sample
View file

@ -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"
}

5
log.py
View file

@ -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)

16
main.py
View file

@ -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)