From c5e5dca000b186e00aa628b628c95298ac16fcbb Mon Sep 17 00:00:00 2001 From: William P Date: Thu, 13 Jun 2024 17:45:12 +0000 Subject: [PATCH] cleanup and slight refactor --- db.py | 25 +++++++------- main.py | 39 ++++++++++++++-------- testing.py | 29 +++++++++++++++++ totp.py | 2 +- vrcapi.py | 96 +++++++++++++++++++++--------------------------------- 5 files changed, 105 insertions(+), 86 deletions(-) create mode 100644 testing.py diff --git a/db.py b/db.py index 89043cc..98bf32e 100644 --- a/db.py +++ b/db.py @@ -1,15 +1,8 @@ -from sqlalchemy import create_engine, Table, Column, MetaData, Integer, String, Boolean, DateTime, ForeignKey -from sqlalchemy.ext.declarative import declarative_base -from sqlalchemy.orm import sessionmaker, relationship, mapped_column import os +from sqlalchemy import create_engine, Column, Integer, String, DateTime, Boolean +from sqlalchemy.ext.declarative import declarative_base +from sqlalchemy.orm import sessionmaker -# Set up SQLite database -if os.getenv("SQLITE_FILE") is not None: - sqlite_file = os.getenv("SQLITE_FILE") -else: - raise Exception("SQLite File not defined!") -DATABASE_URI = f'sqlite:///{sqlite_file}' -engine = create_engine(DATABASE_URI) Base = declarative_base() class DBGroupInstance(Base): @@ -19,10 +12,14 @@ class DBGroupInstance(Base): instance_id = Column(String, nullable=False) world_name = Column(String, nullable=False) member_count = Column(Integer, nullable=False) - -# Create tables if they don't exist -Base.metadata.create_all(engine) + delivered = Column(Boolean, nullable=False) def get_session(): + sqlite_file = os.getenv("SQLITE_FILE") + if not sqlite_file: + raise Exception("SQLite File not defined!") + database_uri = f'sqlite:///{sqlite_file}' + engine = create_engine(database_uri) + Base.metadata.create_all(engine) Session = sessionmaker(bind=engine) - return Session() \ No newline at end of file + return Session() diff --git a/main.py b/main.py index b963c49..09fc366 100644 --- a/main.py +++ b/main.py @@ -1,22 +1,35 @@ -from vrcapi import getGroupInstances, fakeGroupInstances +import os +import logging +from vrcapi import get_group_instances +from testing import fake_group_instances from datetime import datetime from db import DBGroupInstance, get_session +log_level = os.getenv("LOG_LEVEL") +logging.basicConfig(level=log_level) session = get_session() +# Commit the group instance data in the DB +def store_instance_data(instance): + try: + db_instance = DBGroupInstance( + timestamp=datetime.now(), + instance_id=instance.instance_id, + world_name=instance.world.name, + member_count=instance.member_count, + delivered=False # Discord delivery, always false until separate job marks it as True + ) + session.add(db_instance) + session.commit() + logging.info("Instance log added") + except Exception as e: + logging.error(f"Exception while adding instance log: {e}") + session.rollback() + if __name__ == "__main__": - group_instances = getGroupInstances() + group_instances = get_group_instances() # Call the VRChat API and pull group instance data if not group_instances: - print("No group instances were found! Nothing to do...") + logging.info("No group instances were found! Nothing to do...") # Since this will run in a cron job, the program ends here else: for instance in group_instances: - try: - db_instance = DBGroupInstance(timestamp=datetime.now(), - instance_id=instance.instance_id, - world_name=instance.world.name, - member_count=instance.member_count) - session.add(db_instance) - session.commit() - print("Instance log added") - except Exception as e: - print(f"Exception {e}") \ No newline at end of file + store_instance_data(instance) # Send instance data to the DB \ No newline at end of file diff --git a/testing.py b/testing.py new file mode 100644 index 0000000..6e8e744 --- /dev/null +++ b/testing.py @@ -0,0 +1,29 @@ +from vrchatapi.models.group_instance import GroupInstance +from vrchatapi.models.world import World +from random import randint +from datetime import datetime + +def fake_group_instances(): + group_instances = [] + # For testing purposes + + # Worlds API object that is valid enough to not cause errors in testing + fakeWorld = World( + "usr_123TEST-789-abcd-efghi-aaa69420", + "McLovin", + 90, + 50, + datetime.now(), + "A test world, it does not exist. No VRC API's were harmed", + 0, False, 0, "world_AAAA-TEST-49494-69420", + "http://null.example/hello.jpg", [], "4/20/2020", + "Fake world", None, 0, "vrchat", 0, None, 0, 0, "9/11/2021", "public", ["nothing", "testing"], + "http://testing.localhost/thisisnotreal.png", [], datetime.now(), 5, 90, []) + + # Fake instance + fakeInstance1 = GroupInstance("TESTTESTgrp_69420_aaaaaaaaa", + "test", fakeWorld, randint(1, 50)) + + group_instances.append(fakeInstance1) + + return group_instances \ No newline at end of file diff --git a/totp.py b/totp.py index 9dd6ef4..5c7761c 100644 --- a/totp.py +++ b/totp.py @@ -1,7 +1,7 @@ import pyotp import os -def getTOTP(): +def get_totp(): totp_key = os.getenv("TOTP_KEY") totp = pyotp.TOTP(totp_key) return totp.now() \ No newline at end of file diff --git a/vrcapi.py b/vrcapi.py index 107d990..88a15be 100644 --- a/vrcapi.py +++ b/vrcapi.py @@ -1,41 +1,13 @@ +import os +import logging import vrchatapi -from vrchatapi.api import authentication_api +from vrchatapi.api import authentication_api, groups_api from vrchatapi.exceptions import UnauthorizedException, ApiException from vrchatapi.models.two_factor_auth_code import TwoFactorAuthCode from vrchatapi.models.two_factor_email_code import TwoFactorEmailCode -from vrchatapi.models.group_instance import GroupInstance -from vrchatapi.models.world import World -from totp import getTOTP -import os -from random import randint -from datetime import datetime +from totp import get_totp -def fakeGroupInstances(): - group_instances = [] - # For testing purposes - - # Worlds API object that is valid enough to not cause errors in testing - fakeWorld = World( - "usr_123TEST-789-abcd-efghi-aaa69420", - "McLovin", - 90, - 50, - datetime.now(), - "A test world, it does not exist. No VRC API's were harmed", - 0, False, 0, "world_AAAA-TEST-49494-69420", - "http://null.example/hello.jpg", [], "4/20/2020", - "Fake world", None, 0, "vrchat", 0, None, 0, 0, "9/11/2021", "public", ["nothing", "testing"], - "http://testing.localhost/thisisnotreal.png", [], datetime.now(), 5, 90, []) - - # Fake instance - fakeInstance1 = GroupInstance("TESTTESTgrp_69420_aaaaaaaaa", - "test", fakeWorld, randint(1, 50)) - - group_instances.append(fakeInstance1) - - return group_instances - -def getGroupInstances(): +def get_group_instances(): group_instances = [] configuration = vrchatapi.Configuration( @@ -47,33 +19,15 @@ def getGroupInstances(): group_id = os.getenv("GROUP_ID") with vrchatapi.ApiClient(configuration) as api_client: - # Set our User-Agent as per VRChat Usage Policy + # Set our User-Agent as per VRChat Usage Policy and call the Auth API api_client.user_agent = "GroupInstanceLogger/0.1alpha me@williamtpeebles.com" - - # Instantiate instances of API classes auth_api = authentication_api.AuthenticationApi(api_client) try: # Calling getCurrentUser on Authentication API logs you in if the user isn't already logged in. current_user = auth_api.get_current_user() except UnauthorizedException as e: - if e.status == 200: - if "Email 2 Factor Authentication" in e.reason: - # Calling email verify2fa if the account has 2FA disabled, give warning that automatic runs cannot be done - print("IMPORTANT: Your VRChat Account is only enabled for **email** based 2-factor authentication! " - "You will not be able to automate runs of GroupInstanceLogger\n" - "Set up your account to use TOTP-based 2FA so the application can log you in automatically\n\n") - auth_api.verify2_fa_email_code(two_factor_email_code=TwoFactorEmailCode(input("Email 2FA Code: "))) - elif "2 Factor Authentication" in e.reason: - # Check if TOTP code is defined in env variables - if os.getenv("TOTP_KEY") is not None: - print("Attempting to authenticate with defined TOTP key") - auth_api.verify2_fa(two_factor_auth_code=TwoFactorAuthCode(getTOTP())) - else: - # TOTP prompt if not configured - # Calling verify2fa if the account has 2FA enabled - auth_api.verify2_fa(two_factor_auth_code=TwoFactorAuthCode(input("Enter 2FA Code: "))) - + handle_auth_exception(auth_api, e) current_user = auth_api.get_current_user() # Call groups APIs @@ -86,9 +40,35 @@ def getGroupInstances(): return group_instances except ApiException as e: print("Exception when calling GroupsApi->get_group: %s\n" % e) - else: - print("Exception when calling API: %s\n", e) - except vrchatapi.ApiException as e: - print("Exception when calling API: %s\n", e) + logging.info(f"Logged in as: {current_user.display_name}") - print("Logged in as:", current_user.display_name) + try: + api_response = groups_api.GroupsApi(api_client),get_group_instances(group_id) + group_instances.extend(api_response) + except ApiException as e: + logging.error(f"Exception when calling GroupsApi->get_group_instances: {e}") + except Exception as e: + logging.error(f"Unexpected Error: {e}") + return group_instances + + + +def handle_auth_exception(auth_api, e): + if e.status == 200: + if "Email 2 Factor Authentication" in e.reason: + # Calling email verify2fa if the account has 2FA disabled, give warning that automatic runs cannot be done + logging.warning("IMPORTANT: Your VRChat Account is only enabled for **email** based 2-factor authentication! " + "You will not be able to automate runs of GroupInstanceLogger\n" + "Set up your account to use TOTP-based 2FA so the application can log you in automatically\n\n") + auth_api.verify2_fa_email_code(two_factor_email_code=TwoFactorEmailCode(input("Email 2FA Code: "))) + elif "2 Factor Authentication" in e.reason: + # Check if TOTP code is defined in env variables + totp_key = os.getenv("TOTP_KEY") + if totp_key: + logging.info("Attempting to authenticate with defined TOTP key") + auth_api.verify2_fa(two_factor_auth_code=TwoFactorAuthCode(get_totp())) + else: + logging.info("No TOTP key defined, prompting manual input of 2FA code") + auth_api.verify2_fa(two_factor_auth_code=TwoFactorAuthCode(input("Enter 2FA Code: "))) + else: + logging.error(f"UnauthorizedException: {e}") \ No newline at end of file