mirror of
https://github.com/hibobmaster/matrix_chatgpt_bot.git
synced 2024-10-01 05:35:36 -04:00
feat: import E2E room keys
This commit is contained in:
hibobmaster
parent
cb0c913547
commit
7fecffad54
@ -11,3 +11,5 @@ 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
|
||||
IMPORT_KEYS_PATH="element-keys.txt" # Optional
|
||||
IMPORT_KEYS_PASSWORD="xxxxxxx" # Optional
|
||||
41
bot.py
41
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:
|
||||
|
||||
|
||||
@ -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
5
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)
|
||||
|
||||
16
main.py
16
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)
|
||||
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user