Fixed: encryption issues, file pointer inconsistency; Added a new parameter for get_api_key, changed attributes from instance attributes to variables accessed via __getattr__, proper atexit handling

steam/account.py

@@ -17,6 +17,8 @@ import sys
 import json
 import base64
 import re
+import atexit
+from base64 import b64decode
 
 from cryptography.fernet import Fernet
 from cryptography.hazmat.primitives import hashes
@@ -43,6 +45,7 @@ class SteamAccount(object):
     web_api             = None
     authenticator       = None
 
+    _credentials        = { }
     _path               = None
     _file               = None
     _fernet_key         = None
@@ -52,12 +55,29 @@ class SteamAccount(object):
 
 
     def __init__(self, username, password):
+        atexit.register(self.__cleanup__)
         self.username = username
         self.password = password
         if not self._setup():
             raise SteamAccountException('Could not access account.')
 
+    def __getattr__(self, key):
+        if key not in self._credentials.keys():
+            raise AttributeError("No %s attribute" % repr(key))
+        return self._credentials.get(key)
+
     def __del__(self):
+        self.__save_account_credentials__()
+
+    def __cleanup__(self):
+        self.__save_account_credentials__()
+
+    def __save_account_credentials__(self):
+        """
+        Make sure it doesnt write the file if the memory has been cleared
+        """
+        if self._count_account_credentials() <= 2:
+            return
         try:
             self._update_credential_file()
         except TypeError:
@@ -65,29 +85,33 @@ class SteamAccount(object):
             Ignore TypeError exception when destructor gets called after the memory has been cleared
             """
             pass
-        self._file.close()
+        except ValueError:
+            """
+            Ignore ValueError exception when the file could not be written
+            """
+            pass
 
     def set_account_property(self, property, value):
-        setattr(self, property, value)
+        self._credentials[property] = value
         self._update_credential_file()
 
     def del_account_property(self, property):
-        delattr(self, property)
+        del self._credentials[property]
         self._update_credential_file()
 
     def check_account_property(self, property):
-        return hasattr(self, property)
+        return property in self._credentials
 
     @property
     def login_code(self):
-        if self.authenticator:
+        if self.authenticator and hasattr(self.authenticator, 'shared_secret'):
             return self.authenticator.get_code()
         elif hasattr(self, 'shared_secret'):
-            return steam.guard.generate_twofactor_code(self.shared_secret)
+            return steam.guard.generate_twofactor_code(b64decode(self.shared_secret))
         else:
             raise SharedSecretNotSet('Add shared_secret to this instance to generate login codes')
 
-    def get_api_key(self, retrieve_if_missing=True):
+    def get_api_key(self, retrieve_if_missing=False, hostname_for_retrieving='localhost.com'):
         if self.check_account_property('apikey'):
             return self.apikey
 
@@ -95,14 +119,14 @@ class SteamAccount(object):
             if not self._has_web_session():
                 self._spawn_web_session()
 
-            api_key = self.retrieve_api_key()
+            api_key = self.retrieve_api_key(hostname_for_retrieving)
             self.set_account_property('apikey', api_key)
             self._spawn_web_api()
             return api_key
         else:
             raise APIKeyException('Could not return the apikey. The apikey is not set as account property and retrieve_if_missing is not allowed.')
 
-    def retrieve_api_key(self):
+    def retrieve_api_key(self, hostname_for_retrieving='localhost.com'):
         if not self._has_web_session():
             raise APIKeyException('A web session is required to retrieve the api key.')
 
@@ -117,7 +141,7 @@ class SteamAccount(object):
                 session_id = regex_result.group(1)
 
                 data = {
-                    'domain': 'localhost.com',
+                    'domain': hostname_for_retrieving,
                     'agreeToTerms': 'agreed',
                     'submit': 'Register',
                     'sessionid': session_id
@@ -138,16 +162,16 @@ class SteamAccount(object):
         self._generate_fernet_key()
         self._spawn_fernet_suite()
         self._path = '%s/%s' % (BASE_LOCATION, self.username)
-        if not os.path.isfile(self._path):
+        if not os.path.exists(self._path) or not os.path.isfile(self._path):
             self._create_credential_file()
         else:
-            self._file = open(self._path, 'r+', 0)
             credentials = self._parse_credential_file()
             for key, value in credentials.iteritems():
-                setattr(self, key, value)
+                self._credentials[key] = value
 
             if self.check_account_property('shared_secret'):
                 self._spawn_authenticator()
+
             else:
                 self.authenticator = steam.guard.SteamAuthenticator()
 
@@ -156,37 +180,52 @@ class SteamAccount(object):
         return True
 
     def _create_credential_file(self):
-        open(self._path, 'a').close()
+        open(self._path, 'w+').close()
-        self._file = open(self._path, 'r+', 0)
+        self._spawn_file_pointer()
+
         data = json.dumps({
             'username': self.username,
             'password': self.password
         })
+
         token = self._fernet_suite.encrypt(data)
         self._file.write(token)
+        self._file.close()
 
     def _parse_credential_file(self):
-        self._file.seek(0)
+        self._spawn_file_pointer()
+
         token = self._file.read()
-        self._file.seek(0)
         data = json.loads(self._fernet_suite.decrypt(token))
+        self._file.close()
         return data
 
     def _update_credential_file(self):
+        self._spawn_file_pointer()
+
         credentials = self._gather_credentials()
         data = json.dumps(credentials)
         token = self._fernet_suite.encrypt(data)
+
         self._file.truncate()
         self._file.write(token)
+        self._file.close()
 
     def _gather_credentials(self):
         data = { }
-        names = dir(self)
+        names = self._credentials
         for name in names:
             if name in ACCOUNT_ATTRIBUTES:
-                data.__setitem__(name, getattr(self, name))
+                data.__setitem__(name, self._credentials[name])
         return data
 
+    def _count_account_credentials(self):
+        count = 0
+        for attr in ACCOUNT_ATTRIBUTES:
+            if self.check_account_property(attr):
+                count += 1
+        return count
+
     def _generate_fernet_key(self):
         digest = hashes.Hash(hashes.SHA256(), backend=default_backend())
         digest.update(bytes(self.password))
@@ -195,75 +234,59 @@ class SteamAccount(object):
     def _spawn_fernet_suite(self):
         self._fernet_suite = Fernet(self._fernet_key)
 
+    def _spawn_file_pointer(self):
+        self._file = open(self._path, 'r+', 0)
+
     def _spawn_web_api(self):
         self.web_api = steam.webapi.WebAPI(self.apikey)
 
     def _spawn_authenticator(self):
         secrets = {
-                    'identity_secret': getattr(self, 'identity_secret'),
+            'identity_secret':  self._credentials.get('identity_secret'),
-                    'shared_secret': getattr(self, 'shared_secret'),
+            'shared_secret':    self._credentials.get('shared_secret'),
-                    'secret_1': getattr(self, 'secret_1'),
+            'secret_1':         self._credentials.get('secret_1'),
-                    'revocation_code': getattr(self, 'revocation_code'),
+            'revocation_code':  self._credentials.get('revocation_code'),
-                }
+            'deviceid':         self._credentials.get('deviceid')
+        }
         self.authenticator = steam.guard.SteamAuthenticator(secrets)
         self._spawn_mobile_session()
         self.authenticator.medium = self._mobile_auth
 
     def _has_session(self):
-        if self._has_web_session() or self._has_mobile_session():
+        return True if self._has_web_session() or self._has_mobile_session() else False
-            return True
-        return False
 
     def _has_web_session(self):
-        if isinstance(self._web_auth, steam.webauth.WebAuth):
+        return isinstance(self._web_auth, steam.webauth.WebAuth)
-            return True
-        return False
 
     def _has_mobile_session(self):
-        if isinstance(self._mobile_auth, steam.webauth.MobileWebAuth):
+        return isinstance(self._mobile_auth, steam.webauth.MobileWebAuth)
-            return True
-        return False
 
-    def _spawn_web_session(self, captcha='', email_code='', twofactor_code=''):
+    def _spawn_web_session(self):
         self._web_auth = steam.webauth.WebAuth(self.username, self.password)
-        self._login_web_session(self._web_auth, captcha, email_code, twofactor_code)
+        self._login_web_session(self._web_auth)
         self.web_session = self._web_auth.session
 
-    def _spawn_mobile_session(self, captcha='', email_code='', twofactor_code=''):
+    def _spawn_mobile_session(self):
         self._mobile_auth = steam.webauth.MobileWebAuth(self.username, self.password)
-        self._login_web_session(self._mobile_auth, captcha, email_code, twofactor_code)
+        self._login_web_session(self._mobile_auth)
         self.mobile_session = self._mobile_auth.session
 
-    def _login_web_session(self, web_auth, captcha='', email_code='', twofactor_code=''):
+    def _login_web_session(self, web_auth):
         if not isinstance(web_auth, steam.webauth.WebAuth) and not isinstance(web_auth, steam.webauth.MobileWebAuth):
             raise WebAuthNotComplete('Please supply a valid WebAuth or MobileWebAuth session')
 
         try:
-            web_auth.login()
+            twofactor_code = self.login_code
-
+        except SharedSecretNotSet:
-        except steam.webauth.CaptchaRequired:
+            twofactor_code = ''
-            if not captcha:
+
-                raise CaptchaNotProvided('The steam login captcha is required for logging in, but was not provided.')
+        web_auth.login(twofactor_code=twofactor_code)
-            web_auth.login(captcha=captcha)
-
-        except steam.webauth.EmailCodeRequired:
-            if not email_code:
-                raise EMailCodeNotProvided('The email code is required for logging in, but was not provided.')
-            web_auth.login(email_code=email_code)
-
-        except steam.webauth.TwoFactorCodeRequired:
-            if not twofactor_code:
-                try:
-                    twofactor_code = self.login_code
-                except SharedSecretNotSet:
-                    raise TwoFACodeNotProvided('The twofactor code is required for logging in, but was not provided.')
-            web_auth.login(twofactor_code=twofactor_code)
 
         if web_auth.complete:
             if not hasattr(self, 'steamid'):
                 self.set_account_property('steamid', web_auth.steam_id)
 
-            if isinstance(web_auth, steam.webauth.MobileWebAuth) and not hasattr(self, 'oauth_token'):
+            if isinstance(web_auth, steam.webauth.MobileWebAuth):
                 self.set_account_property('oauth_token', web_auth.oauth_token)
         else:
             raise WebAuthNotComplete('The web authentication could not be completed.')
@@ -286,20 +309,5 @@ class WebException(SteamAccountException):
 class APIKeyException(SteamAccountException):
     pass
 
-class MobileAuthenticatorException(SteamAccountException):
-    pass
-
 class ParameterNotProvidedException(SteamAccountException):
-    pass
-
-class CaptchaNotProvided(ParameterNotProvidedException):
-    pass
-
-class EMailCodeNotProvided(ParameterNotProvidedException):
-    pass
-
-class TwoFACodeNotProvided(ParameterNotProvidedException):
-    pass
-
-class SMSCodeNotProvided(ParameterNotProvidedException):
     pass