Handle unauthorized response and raise UnauthorizeError

diyan · Aug 20, 2014 · 7ab74a4 · 7ab74a4
1 parent df04945
commit 7ab74a4
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 6 deletions.
diff --git a/winrm/exceptions.py b/winrm/exceptions.py
@@ -1,3 +1,6 @@
+import re
+
+
 class WinRMWebServiceError(Exception):
     """Generic WinRM SOAP Error"""
     pass
@@ -15,4 +18,24 @@ class WinRMWSManFault(Exception):
 
 class WinRMTransportError(Exception):
     """"Transport-level error"""
-    pass
+    code = 500
+
+    def __init__(self, transport, message):
+        self.transport = transport
+        self.message = message
+
+    def __str__(self):
+        return '{0} {1}. {2}'.format(self.code, re.sub('Error$', '', self.__class__.__name__), self.message)
+
+    def __repr__(self):
+        return "{0}(code={1}, transport='{2}', message='{3}')".format(
+            self.__class__.__name__, self.code, self.transport, self.message)
+
+
+class UnauthorizedError(WinRMTransportError):
+    """Raise if the user is not authorized"""
+    code = 401
+
+
+class TimeoutError(WinRMTransportError):
+    pass
diff --git a/winrm/transport.py b/winrm/transport.py
@@ -1,6 +1,8 @@
 import sys
 import base64
-from winrm.exceptions import WinRMTransportError
+
+from winrm.exceptions import WinRMTransportError, UnauthorizedError
+
 
 HAVE_KERBEROS=False
 try:
@@ -53,8 +55,8 @@ def __init__(self, endpoint, username='', password='', disable_sspi=True, basic_
         if basic_auth_only:
             self.basic_auth_only()
 
-        self._headers = {'Content-Type' : 'application/soap+xml;charset=UTF-8',
-                         'User-Agent' : 'Python WinRM client'}
+        self._headers = {'Content-Type': 'application/soap+xml;charset=UTF-8',
+                         'User-Agent': 'Python WinRM client'}
 
     def _setup_opener(self):
         password_manager = HTTPPasswordMgrWithDefaultRealm()
@@ -84,17 +86,20 @@ def send_message(self, message):
             #return doc
             #return doc
         except HTTPError as ex:
+            if ex.code == 401:
+                raise UnauthorizedError(transport='plaintext', message=ex.msg)
             response_text = ex.read()
             # Per http://msdn.microsoft.com/en-us/library/cc251676.aspx rule 3,
             # should handle this 500 error and retry receiving command output.
             if 'http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Receive' in message and 'Code="2150858793"' in response_text:
+                # TODO raise TimeoutError here instead of just return text
                 return response_text
             error_message = 'Bad HTTP response returned from server. Code {0}'.format(ex.code)
             if ex.msg:
                 error_message += ', {0}'.format(ex.msg)
-            raise WinRMTransportError(error_message)
+            raise WinRMTransportError('http', error_message)
         except URLError as ex:
-            raise WinRMTransportError(ex.reason)
+            raise WinRMTransportError('http', ex.reason)
 
 
 class HTTPSClientAuthHandler(HTTPSHandler):