class Attachment:
def __init__(self, part):
encoding = part.encoding or "utf-8"
self.headers = CaseInsensitiveDict({
k.decode(encoding): v.decode(encoding)
for k, v in part.headers.items()
})
self.content_type = self.headers.get("Content-Type", None)
self.content_id = self.headers.get("Content-ID", None)
self.content_location = self.headers.get("Content-Location", None)
self._part = part
def __repr__(self):
return "<Attachment(%r, %r)>" % (self.content_id, self.content_type)
@cached_property
def content(self):
"""Return the content of the attachment
:rtype: bytes or str
"""
encoding = self.headers.get("Content-Transfer-Encoding", None)
content = self._part.content
if encoding == "base64":
return base64.b64decode(content)
elif encoding == "binary":
return content.strip(b"\r\n")
else:
return content
def string_to_sign(self, request):
"""
Generates the string we need to sign on.
Params:
- request The request object
Returns
String ready to be signed on
"""
# We'll use case insensitive dict to store the headers
h = CaseInsensitiveDict()
# Add the hearders
h.update(request.headers)
# If we have an 'x-amz-date' header,
# we'll try to use it instead of the date
if b'x-amz-date' in h or 'x-amz-date' in h:
date = ''
else:
# No x-amz-header, we'll generate a date
date = h.get('Date') or self._get_date()
# Set the date header
request.headers['Date'] = date
# A fix for the content type header extraction in python 3
# This have to be done because requests will try to set
# application/www-url-encoded header if we pass bytes as the content,
# and the content-type is set with a key that is b'Content-Type' and
# not 'Content-Type'
content_type = ''
if b'Content-Type' in request.headers:
# Fix content type
content_type = h.get(b'Content-Type')
del request.headers[b'Content-Type']
request.headers['Content-Type'] = content_type
# The string we're about to generate
# There's more information about it here:
# http://docs.aws.amazon.com/AmazonS3/latest/dev/
# RESTAuthentication.html#ConstructingTheAuthenticationHeader
msg = [
# HTTP Method
request.method,
# MD5 If provided
h.get(b'Content-MD5', '') or h.get('Content-MD5', ''),
# Content type if provided
content_type or h.get('Content-Type', ''),
# Date
date,
# Canonicalized special amazon headers and resource uri
self._get_canonicalized_amz_headers(h) +
self._get_canonicalized_resource(request)
]
# join with a newline and return
return '\n'.join(msg)
def response(self):
if self._response is None and self.request is not None:
request = self.request
existing = None
if self.use_cache:
existing = self.context.get_tag(self.request_id)
if existing is not None:
headers = CaseInsensitiveDict(existing.get('headers'))
last_modified = headers.get('last-modified')
if last_modified:
request.headers['If-Modified-Since'] = last_modified
etag = headers.get('etag')
if etag:
request.headers['If-None-Match'] = etag
self._rate_limit(request.url)
session = self.http.session
prepared = session.prepare_request(request)
response = session.send(prepared,
stream=True,
verify=False,
allow_redirects=self.allow_redirects)
if existing is not None and response.status_code == 304:
self.context.log.info("Using cached HTTP response: %s",
response.url)
self.apply_data(existing)
else:
self._response = response
# update the serialised session with cookies etc.
self.http.save()
return self._response
class Attachment(object):
def __init__(self, part):
encoding = part.encoding or 'utf-8'
self.headers = CaseInsensitiveDict({
k.decode(encoding): v.decode(encoding)
for k, v in part.headers.items()
})
self.content_type = self.headers.get('Content-Type', None)
self.content_id = self.headers.get('Content-ID', None)
self.content_location = self.headers.get('Content-Location', None)
self._part = part
def __repr__(self):
return '<Attachment(%r, %r)>' % (self.content_id, self.content_type)
@cached_property
def content(self):
"""Return the content of the attachment
:rtype: bytes or str
"""
encoding = self.headers.get('Content-Transfer-Encoding', None)
content = self._part.content
if encoding == 'base64':
return base64.b64decode(content)
elif encoding == 'binary':
return content.strip(b'\r\n')
else:
return content
class Attachment(object):
def __init__(self, part):
self.headers = CaseInsensitiveDict({
k.decode(part.encoding): v.decode(part.encoding)
for k, v in part.headers.items()
})
self.content_type = self.headers.get('Content-Type', None)
self.content_id = self.headers.get('Content-ID', None)
self.content_location = self.headers.get('Content-Location', None)
self._part = part
def __repr__(self):
return '<Attachment(%r, %r)>' % (self.content_id, self.content_type)
@cached_property
def content(self):
encoding = self.headers.get('Content-Transfer-Encoding', None)
content = self._part.content
if encoding == 'base64':
return base64.b64decode(content)
elif encoding == 'binary':
return content
else:
return content
def string_to_sign(self, request):
"""
Generates the string we need to sign on.
Params:
- request The request object
Returns
String ready to be signed on
"""
# We'll use case insensitive dict to store the headers
h = CaseInsensitiveDict()
# Add the hearders
h.update(request.headers)
# If we have an 'x-amz-date' header,
# we'll try to use it instead of the date
if b'x-amz-date' in h or 'x-amz-date' in h:
date = ''
else:
# No x-amz-header, we'll generate a date
date = h.get('Date') or self._get_date()
# Set the date header
request.headers['Date'] = date
# A fix for the content type header extraction in python 3
# This have to be done because requests will try to set
# application/www-url-encoded header if we pass bytes as the content,
# and the content-type is set with a key that is b'Content-Type' and
# not 'Content-Type'
content_type = ''
if b'Content-Type' in request.headers:
# Fix content type
content_type = h.get(b'Content-Type')
del request.headers[b'Content-Type']
request.headers['Content-Type'] = content_type
# The string we're about to generate
# There's more information about it here:
# http://docs.aws.amazon.com/AmazonS3/latest/dev/
# RESTAuthentication.html#ConstructingTheAuthenticationHeader
msg = [
# HTTP Method
request.method,
# MD5 If provided
h.get(b'Content-MD5', '') or h.get('Content-MD5', ''),
# Content type if provided
content_type or h.get('Content-Type', ''),
# Date
date,
# Canonicalized special amazon headers and resource uri
self._get_canonicalized_amz_headers(h) +
self._get_canonicalized_resource(request)
]
# join with a newline and return
return '\n'.join(msg)
def string_to_sign(self, request):
h = CaseInsensitiveDict()
h.update(request.headers)
# Try to use
if b'x-amz-date' in h or 'x-amz-date' in h:
date = ''
else:
date = h.get('Date') or self._get_date()
request.headers['Date'] = date
# Set the date header
request.headers['Date'] = date
# A fix for the content type header extraction in python 3
# This have to be done because requests will try to set application/www-url-encoded herader
# if we pass bytes as the content, and the content-type is set with a key that is b'Content-Type' and not
# 'Content-Type'
content_type = ''
if b'Content-Type' in request.headers:
# Fix content type
content_type = h.get(b'Content-Type')
del request.headers[b'Content-Type']
request.headers['Content-Type'] = content_type
msg = [
request.method,
h.get(b'Content-MD5', '') or h.get('Content-MD5', ''),
content_type or h.get('Content-Type', ''),
date,
self._get_canonicalized_amz_headers(h) + self._get_canonicalized_resource(request)
]
return '\n'.join(msg)
def test_get(self):
cid = CaseInsensitiveDict()
cid['spam'] = 'oneval'
cid['SPAM'] = 'blueval'
assert cid.get('spam') == 'blueval'
assert cid.get('SPAM') == 'blueval'
assert cid.get('sPam') == 'blueval'
assert cid.get('notspam', 'default') == 'default'
def test_get(self):
cid = CaseInsensitiveDict()
cid['spam'] = 'oneval'
cid['SPAM'] = 'blueval'
self.assertEqual(cid.get('spam'), 'blueval')
self.assertEqual(cid.get('SPAM'), 'blueval')
self.assertEqual(cid.get('sPam'), 'blueval')
self.assertEqual(cid.get('notspam', 'default'), 'default')
def test_get(self):
cid = CaseInsensitiveDict()
cid['spam'] = 'oneval'
cid['SPAM'] = 'blueval'
assert cid.get('spam') == 'blueval'
assert cid.get('SPAM') == 'blueval'
assert cid.get('sPam') == 'blueval'
assert cid.get('notspam', 'default') == 'default'
def test_get(self):
cid = CaseInsensitiveDict()
cid["spam"] = "oneval"
cid["SPAM"] = "blueval"
assert cid.get("spam") == "blueval"
assert cid.get("SPAM") == "blueval"
assert cid.get("sPam") == "blueval"
assert cid.get("notspam", "default") == "default"
def test_get(self):
cid = CaseInsensitiveDict()
cid["spam"] = "oneval"
cid["SPAM"] = "blueval"
self.assertEqual(cid.get("spam"), "blueval")
self.assertEqual(cid.get("SPAM"), "blueval")
self.assertEqual(cid.get("sPam"), "blueval")
self.assertEqual(cid.get("notspam", "default"), "default")
def test_get(self):
cid = CaseInsensitiveDict()
cid['spam'] = 'oneval'
cid['SPAM'] = 'blueval'
self.assertEqual(cid.get('spam'), 'blueval')
self.assertEqual(cid.get('SPAM'), 'blueval')
self.assertEqual(cid.get('sPam'), 'blueval')
self.assertEqual(cid.get('notspam', 'default'), 'default')
def aws_invoke(app,gateway_input,server_name='localhost',server_port='5000',http_protocol='HTTP/1.1',TLS=True,block_headers=True):
headers = CaseInsensitiveDict(gateway_input.get('headers',{}))
requestContext = gateway_input.get('requestContext')
queryStringParameters = gateway_input.get('queryStringParameters',{})
clientIp = headers.get('x-forwarded-for')
if clientIp is None:
clientIp = requestContext.get('identity',{}).get('sourceIp') if requestContext is not None else ''
else:
clientIp = clientIp.split(',')[0]
environ = {
'REQUEST_METHOD': gateway_input.get('httpMethod','GET').upper(),
'SCRIPT_NAME': '',
'PATH_INFO': gateway_input.get('path','/'),
'QUERY_STRING': urlencode(queryStringParameters) if queryStringParameters is not None else '',
'SERVER_NAME': headers.get('host',server_name),
'SERVER_PORT': headers.get('x-forwarded-port',server_port),
'SERVER_PROTOCOL': http_protocol,
'SERVER_SOFTWARE': 'flask-serverless',
'REMOTE_ADDR': clientIp,
'wsgi.version': (1, 0),
'wsgi.url_scheme': headers.get('x-forwarded-proto','https' if TLS else 'http'),
'wsgi.input': None,
'wsgi.errors': sys.stderr,
'wsgi.multiprocess': True,
'wsgi.multithread': False,
'wsgi.run_once': True
}
if environ['REQUEST_METHOD']=='POST' or environ['REQUEST_METHOD']=='PUT':
contentType = headers.get('content-type','application/octet-stream')
parsedContentType = parse_options_header(contentType)
raw = gateway_input.get('body')
if raw is None or gateway_input.get('isBase64Encoded',False):
body = b64decode(raw) if raw is not None else None
else:
body = raw.encode(parsedContentType[1].get('charset','utf-8'))
add_body(environ,body,contentType)
add_headers(environ,headers,block_headers)
response = Response.from_app(app.wsgi_app, environ)
gateway_output = {
'headers' : dict(response.headers),
'statusCode' : response.status_code,
}
compressed = response.headers.get('Content-Encoding')=='gzip'
responseType = parse_options_header(response.headers.get('Content-Type','application/octet-stream'))
if not compressed and ('charset' in responseType[1] or responseType[0] in textTypes or responseType[0][0:5]=='text/'):
gateway_output['body'] = response.data.decode(responseType[1].get('charset','utf-8'))
gateway_output['isBase64Encoded'] = False
else:
gateway_output['body'] = b64encode(response.data).decode('utf-8')
gateway_output['isBase64Encoded'] = True
return gateway_output
def _guess_mime_type(headers: CaseInsensitiveDict) -> Optional[str]:
location = headers.get('location')
if location:
mime_type, _ = mimetypes.guess_type(location)
if mime_type:
return mime_type
# Parse mime type from content-type header, e.g. 'image/jpeg;charset=US-ASCII' -> 'image/jpeg'
mime_type, _ = cgi.parse_header(headers.get('content-type', ''))
return mime_type or None
def __init__(self,
response=None,
status=None,
headers=None,
mimetype=None,
content_type=None,
direct_passthrough=False):
headers = CaseInsensitiveDict(headers) if headers is not None else None
if response is not None and isinstance(
response, BaseResponse) and response.headers is not None:
headers = CaseInsensitiveDict(response.headers)
if headers is None:
headers = CaseInsensitiveDict()
h = headers
h['Access-Control-Allow-Origin'] = headers.get(
'Access-Control-Allow-Origin', '*')
h['Access-Control-Allow-Methods'] = headers.get(
'Access-Control-Allow-Methods',
"GET, PUT, POST, HEAD, OPTIONS, DELETE")
h['Access-Control-Max-Age'] = headers.get('Access-Control-Max-Age',
"21600")
h['Cache-Control'] = headers.get(
'Cache-Control', "no-cache, must-revalidate, no-store")
if 'Access-Control-Allow-Headers' not in headers and len(
headers.keys()) > 0:
h['Access-Control-Allow-Headers'] = ', '.join(iterkeys(headers))
data = None
if response is not None and isinstance(response, string_types):
data = response
response = None
if response is not None and isinstance(response, BaseResponse):
new_response_headers = CaseInsensitiveDict(
response.headers if response.headers is not None else {})
new_response_headers.update(h)
response.headers = new_response_headers
headers = None
data = response.get_data()
else:
headers.update(h)
headers = dict(headers)
super(IppResponse,
self).__init__(response=response,
status=status,
headers=headers,
mimetype=mimetype,
content_type=content_type,
direct_passthrough=direct_passthrough)
if data is not None:
self.set_data(data)
def prepare_response(self, cached):
"""Verify our vary headers match and construct a real urllib3
HTTPResponse object.
"""
# Special case the '*' Vary value as it means we cannot actually
# determine if the cached response is suitable for this request.
if "*" in cached.get("vary", {}):
return
body_raw = cached["response"].pop("body")
headers = CaseInsensitiveDict(data=cached['response']['headers'])
if headers.get('transfer-encoding', '') == 'chunked':
headers.pop('transfer-encoding')
cached['response']['headers'] = headers
try:
body = io.BytesIO(body_raw)
except TypeError:
# This can happen if cachecontrol serialized to v1 format (pickle)
# using Python 2. A Python 2 str(byte string) will be unpickled as
# a Python 3 str (unicode string), which will cause the above to
# fail with:
#
# TypeError: 'str' does not support the buffer interface
body = io.BytesIO(body_raw.encode('utf8'))
return HTTPResponse(
body=body,
preload_content=False,
**cached["response"]
)
def prepare_response(self, request, cached):
"""Verify our vary headers match and construct a real urllib3
HTTPResponse object.
"""
# Special case the '*' Vary value as it means we cannot actually
# determine if the cached response is suitable for this request.
if "*" in cached.get("vary", {}):
return
# Ensure that the Vary headers for the cached response match our
# request
for header, value in cached.get("vary", {}).items():
if request.headers.get(header, None) != value:
return
body_file = cached[u'response'].pop(u'body')
body = open(os.path.join(self.cache.directory, body_file), 'rb')
headers = CaseInsensitiveDict(data=cached['response']['headers'])
if headers.get('transfer-encoding', '') == 'chunked':
headers.pop('transfer-encoding')
cached['response']['headers'] = headers
return HTTPResponse(
body=body,
preload_content=False,
**cached["response"]
)
def call_processor_with_span_context(self, seqid, iprot, oprot):
context = baseplate.make_context_object()
# Allow case-insensitivity for THeader headers
headers = CaseInsensitiveDict(data=iprot.get_headers())
try:
trace_info = _extract_trace_info(headers)
except (KeyError, ValueError):
trace_info = None
edge_payload = headers.get(b"Edge-Request", None)
if edge_context_factory:
edge_context = edge_context_factory.from_upstream(edge_payload)
edge_context.attach_context(context)
else:
# just attach the raw context so it gets passed on
# downstream even if we don't know how to handle it.
context.raw_request_context = edge_payload
baseplate.make_server_span(context,
name=fn_name,
trace_info=trace_info)
context.headers = headers
handler = processor._handler
context_aware_handler = _ContextAwareHandler(
handler, context, logger)
context_aware_processor = processor.__class__(
context_aware_handler)
return processor_fn(context_aware_processor, seqid, iprot, oprot)
class DDWRT(Router):
def __init__(self, conf, hostnames):
self.hostnames = CaseInsensitiveDict()
self.hostnames.update(hostnames)
self.conf = conf
self.auth = self.conf.auth()
def clients(self):
""" Receives all currently logged in users in a wifi network.
:rtype : list
:return: Returns a list of dicts, containing the following keys: mac, ipv4, seen, hostname
"""
clients = self._get_clients_raw()
clients_json = []
for client in clients:
client_hostname_from_router = client[0]
client_ipv4 = client[1].strip()
client_mac = client[2].strip().upper()
client_hostname = self.hostnames.get(client_mac, client_hostname_from_router).strip()
client_connections = int(client[3].strip())
# Clients with less than 20 connections are considered offline
if client_connections < 20:
continue
clients_json.append({
'mac': client_mac,
'ipv4': client_ipv4,
'seen': int(time.time()),
'hostname': client_hostname,
})
logger.debug('The router got us {} clients.'.format(len(clients_json)))
logger.debug(str(clients_json))
return clients_json
def _get_clients_raw(self):
info_page = self.conf.internal()
response = requests.get(info_page, auth=self.auth)
logger.info('Got response from router with code {}.'.format(response.status_code))
return DDWRT._convert_to_clients(response.text) or []
@staticmethod
def _convert_to_clients(router_info_all):
# Split router info in lines and filter empty info
router_info_lines = filter(None, router_info_all.split("\n"))
# Get key / value of router info
router_info_items = dict()
for item in router_info_lines:
key, value = item[1:-1].split("::") # Remove curly braces and split
router_info_items[key.strip()] = value.strip()
# Get client info as a list
arp_table = utils.groupn(router_info_items['arp_table'].replace("'", "").split(","), 4)
dhcp_leases = utils.groupn(router_info_items['dhcp_leases'].replace("'", "").split(","), 5)
return arp_table if (len(arp_table) > 0) else []
def _index_into_idol(self, documents, query):
index_data = ''
for _d in documents:
fields = _d.get('fields', [])
content = _d.get('drecontent', '')
DOCUMENTS = _d.get('content', {}).get('DOCUMENT', [])
for DOC in DOCUMENTS:
for key in DOC:
if key == 'DRECONTENT':
for value in DOC[key]:
content += value
else:
for value in DOC[key]:
fields.append((key, value))
index_data += '\n'.join(
[f"#DREREFERENCE {_d.get('reference')}"] +
[f"#DREFIELD {_f[0]}=\"{_f[1]}\"" for _f in fields] +
[f"#DRECONTENT", f"{content}", "#DREENDDOC\n\n"])
# add to queue
_query = CaseInsensitiveDict(query)
if _query.get('priority', 0) >= 100: # bypass queue
self.post_index_data(_query,
[(None, _query, index_data, len(documents))])
else:
self.add_into_batch_queue(_query, index_data, len(documents))
def prepare_response(self, cached):
"""Verify our vary headers match and construct a real urllib3
HTTPResponse object.
"""
# Special case the '*' Vary value as it means we cannot actually
# determine if the cached response is suitable for this request.
if "*" in cached.get("vary", {}):
return
body_raw = cached["response"].pop("body")
headers = CaseInsensitiveDict(data=cached['response']['headers'])
if headers.get('transfer-encoding', '') == 'chunked':
headers.pop('transfer-encoding')
cached['response']['headers'] = headers
try:
body = io.BytesIO(body_raw)
except TypeError:
# This can happen if cachecontrol serialized to v1 format (pickle)
# using Python 2. A Python 2 str(byte string) will be unpickled as
# a Python 3 str (unicode string), which will cause the above to
# fail with:
#
# TypeError: 'str' does not support the buffer interface
body = io.BytesIO(body_raw.encode('utf8'))
return HTTPResponse(body=body,
preload_content=False,
**cached["response"])
def parser_cookies_by_headers(headers: CaseInsensitiveDict):
cookies = RequestsCookieJar()
cookie_str = headers.get("Cookie")
if cookie_str:
for cookie in cookie_str.strip().split(";"):
key, value = cookie.strip().split("=")
cookies.set(key.strip(), value.strip())
return cookies
def get_file_name(headers: CaseInsensitiveDict) -> str:
file_name = str(uuid.uuid4()) + ".csv"
content_disposition = headers.get("Content-Disposition")
if content_disposition is not None:
if "attachment;" in content_disposition:
fn_directive = content_disposition.split(" ")[1]
file_name = fn_directive.split("=")[1]
return file_name
def build_response(self, req, resp):
response = requests.adapters.HTTPAdapter.build_response(
self, req, resp)
if response.headers.get("content-type") != None:
if response.headers["content-type"].startswith(
"multipart/encrypted"):
_, options_part = response.headers["content-type"].split(
";", 1)
options = CaseInsensitiveDict()
for item in options_part.split(";"):
key, value = item.split("=")
if value[0] == '"' and value[-1] == '"':
value = value[1:-1]
value = value.replace(b'\\\\',
b'\\').replace(b'\\"', b'"')
options[key] = value
if options.get("protocol") is not None and options[
"protocol"] == "application/HTTP-Kerberos-session-encrypted":
boundary = options["boundary"]
encrypted_data = None
re_multipart = r'(?:--' + boundary + r'(?:(?:\r\n)|(?:--(?:\r\n)*)))'
for part in re.split(re_multipart, response.content):
if part == '':
continue
(header_raw, data) = part.split('\r\n', 1)
key, value = map(lambda x: x.strip(),
header_raw.split(":"))
if key.lower(
) == "content-type" and value == "application/HTTP-Kerberos-session-encrypted":
_, orginaltype = map(lambda x: x.strip(),
data.split(":"))
original_values = CaseInsensitiveDict()
for item in orginaltype.split(";"):
subkey, subvalue = item.split("=")
original_values[subkey] = subvalue
if key.lower(
) == "content-type" and value == "application/octet-stream":
encrypted_data = data
con = self.get_connection(req.url, None)
decrypted = HTTPMSKerberosAdapter.krb_dict[
req.url].decrypt(encrypted_data)
response.headers["Content-Type"] = original_values[
"type"] + "; charset=" + original_values["charset"]
response.headers["Content-Length"] = len(decrypted)
response.encoding = requests.utils.get_encoding_from_headers(
response.headers)
response.raw = StringIO.StringIO(decrypted)
response._content_consumed = False
response._content = False
return response
def get_retry_after(response):
"""Get the value of Retry-After in seconds.
:param response: The PipelineResponse object
:type response: ~azure.core.pipeline.PipelineResponse
:return: Value of Retry-After in seconds.
:rtype: float or None
"""
headers = CaseInsensitiveDict(response.http_response.headers)
retry_after = headers.get("retry-after")
if retry_after:
return parse_retry_after(retry_after)
for ms_header in ["retry-after-ms", "x-ms-retry-after-ms"]:
retry_after = headers.get(ms_header)
if retry_after:
parsed_retry_after = parse_retry_after(retry_after)
return parsed_retry_after / 1000.0
return None
class GroupStore:
"""
Simple in-memory store for group info. Attribute names are case-insensitive.
Users can be retrieved by id or groupname
"""
def __init__(self):
self.groups = CaseInsensitiveDict() # key = id
self.names = CaseInsensitiveDict() # key = groupname
def add(self, info):
self.groups[info.get('id')] = info
self.names[info.get('displayname')] = info
return id
def get_by_id(self, id):
return self.groups.get(id, None)
def get_by_name(self, username):
return self.names.get(username, None)
class UserStore:
"""
Simple in-memory store for user info. Attribute names are case-insensitive.
Users can be retrieved by id or username
"""
def __init__(self):
self.users = CaseInsensitiveDict() # key = id
self.names = CaseInsensitiveDict() # key = username
def add(self, user_info):
self._add_default_attributes(user_info)
self.users[user_info.get('id')] = dict(user_info)
self.names[user_info.get('username')] = dict(user_info)
return id
def get_by_id(self, id):
return self.users.get(id, None)
def get_by_name(self, username):
return self.names.get(username, None)
def update_scopes(self, username, scopes):
self.names[username]['consented_scopes'] += ' ' + scopes
def list(self):
"""
Returns a list of dictionaries representing users.
password and consented_scopes attributes are not returned
"""
return [self._copy_user(u[1]) for u in self.users.items()]
def _copy_user(self, user):
d = copy.deepcopy(dict(user))
self._del_default_attributes(d)
return d
def _add_default_attributes(self, user_info):
if 'consented_scopes' not in user_info:
user_info['consented_scopes'] = ''
def _del_default_attributes(self, dictionary):
del dictionary['consented_scopes']
del dictionary['password']
class Response:
def __init__(self, request):
"""Response object corresponding to the request object
"""
self.request = request
self.status_code = None
self.content = bytearray()
self.headers = CaseInsensitiveDict()
def read_status(self):
return int(self.request.file.readline().decode('utf-8').split()[1])
def read_headers(self):
for line in self.request.file:
if line == b'\r\n': # end of headers
break
header = line.decode().strip(
) # remove leading and trailing white spaces
key, value = header.split(':', maxsplit=1)
self.headers[key] = value.strip()
def read_content(self):
file = self.request.file
## Content-length specified:
if self.headers.get('Content-length',
None): # Content-length header exists
self.content = file.read(int(self.headers['Content-length']))
## Chunked transfer specified: chunk represented as hexa string + CRLF
elif self.headers.get('Transfer-Encoding',
None) == 'chunked': # chunked transfer
while True:
content_length = int(file.readline().decode('utf-8').strip(),
16)
if content_length == 0:
break
chunk = file.read(content_length)
self.content.extend(chunk)
file.readline()
file.readline() # skip CRLF
## nothing specified:
else:
self.content = file.read() # read until FIN arrival
def _get_email_headers_from_part(self, part, charset=None):
email_headers = list(part.items())
# TODO: the next 2 ifs can be condensed to use 'or'
if charset is None:
charset = part.get_content_charset()
if charset is None:
charset = 'utf8'
if not email_headers:
return {}
# Convert the header tuple into a dictionary
headers = CaseInsensitiveDict()
try:
[
headers.update({x[0]: self._get_string(x[1], charset)})
for x in email_headers
]
except Exception as e:
error_code, error_msg = self._get_error_message_from_exception(e)
err = "Error occurred while converting the header tuple into a dictionary"
self._base_connector.debug_print("{}. {}. {}".format(
err, error_code, error_msg))
# Handle received separately
try:
received_headers = [
self._get_string(x[1], charset) for x in email_headers
if x[0].lower() == 'received'
]
except Exception as e:
error_code, error_msg = self._get_error_message_from_exception(e)
err = "Error occurred while handling the received header tuple separately"
self._base_connector.debug_print("{}. {}. {}".format(
err, error_code, error_msg))
if received_headers:
headers['Received'] = received_headers
# handle the subject string, if required add a new key
subject = headers.get('Subject')
if subject:
try:
headers['decodedSubject'] = str(
make_header(decode_header(subject)))
except Exception:
headers['decodedSubject'] = self._decode_uni_string(
subject, subject)
return dict(headers)
def _cache_ratelimit_headers(self,
headers: CaseInsensitiveDict,
resource: str = CORE_RESOURCE) -> dict:
"""Cache rate limit information from response headers.
:param requests.structures.CaseInsensitiveDict headers:
Headers from response.
:param str resource:
Name of resource to get rate limit for. Either CORE_RESOURCE,
SEARCH_RESOURCE, or GRAPHQL_RESOURCE.
:returns dict:
Dictionary containing remaining rate limit, full rate limit, and
reset time as POSIX timestamp. For more information see
https://developer.github.com/v3/rate_limit/
"""
if not self._ratelimit_cache:
self._ratelimit_cache = {}
if self._has_ratelimit_headers(headers):
self._ratelimit_cache[resource] = {
'limit': headers.get(self.RATELIMIT_LIMIT_HEADER),
'remaining': headers.get(self.RATELIMIT_REMAINING_HEADER),
'reset': headers.get(self.RATELIMIT_RESET_HEADER)
}
class UserStore:
"""
Simple in-memory store for user info. Attribute names are case-insensitive.
Users can be retrieved by id or username
"""
def __init__(self):
self.users = CaseInsensitiveDict()
self.names = CaseInsensitiveDict()
def add(self, info):
self.users[info.get('id')] = info
self.names[info.get('username')] = info
return id
def get_by_id(self, id):
return self.users.get(id, None)
def get_by_name(self, username):
return self.names.get(username, None)
def update_scopes(self, username, scopes):
self.users[username]['consented_scopes'] += ' ' + scopes
def _get_conflicting_active_course_enrollments(requests_by_key,
existing_course_enrollments,
program_uuid, course_key):
"""
Process the list of existing course enrollments together with
the enrollment request list stored in 'requests_by_key'. Detect
whether we have conflicting ACTIVE ProgramCourseEnrollment entries.
When detected, log about it and return the conflicting entry with
duplicated status.
Arguments:
requests_by_key (dict)
existing_course_enrollments (queryset[ProgramCourseEnrollment]),
program_uuid (UUID|str),
course_key (str)
Returns:
results (dict) with detected conflict entry, or empty dict.
"""
conflicted_by_user_key = CaseInsensitiveDict()
requested_statuses_by_user_key = CaseInsensitiveDict({
key: request.get('status')
for key, request in requests_by_key.items()
})
for existing_enrollment in existing_course_enrollments:
external_user_key = existing_enrollment.program_enrollment.external_user_key
requested_status = requested_statuses_by_user_key.get(
existing_enrollment.program_enrollment.external_user_key)
if (requested_status
and requested_status == ProgramCourseEnrollmentStatuses.ACTIVE
and existing_enrollment.is_active
and str(existing_enrollment.program_enrollment.program_uuid) !=
str(program_uuid)):
logger.error(
'Detected conflicting active ProgramCourseEnrollment. This is happening on'
' The program_uuid [{}] with course_key [{}] for external_user_key [{}]'
.format(program_uuid, course_key, external_user_key))
conflicted_by_user_key[
external_user_key] = ProgramCourseOpStatuses.CONFLICT
return conflicted_by_user_key
class Response:
"""The object that each response must be packed into before sending. Same
reason as the Request object. """
__slots__ = ['version', 'status', 'headers', 'body']
def __init__(self):
self.version = HTTP1_1
self.status = 200
self.headers = CID()
self.body = bytes()
def __repr__(self):
return '<Response %s %s>' % (self.status, self.reason)
@property
def reason(self):
"""Third argument in response status line"""
return responses[self.status]
def unpack(self):
"""Preprocess http message (e.g decode)"""
if self.headers.get('content-encoding', '') == 'gzip':
self.body = gzip.decompress(self.body)
def prepare(self, request: Request) -> bytes:
"""Prepare the response for socket transmission"""
if 'gzip' in request.headers.get('accept-encoding', ''):
self.headers['content-encoding'] = 'gzip'
self.body = gzip.compress(self.body)
self.headers['content-length'] = str(len(self.body))
return '{version} {status} {reason}{headers}\r\n\r\n'.format(
version=self.version,
status=self.status,
reason=self.reason,
headers=''.join(
'\r\n%s: %s' % (k, v)
for k, v in self.headers.items())).encode() + self.body
def prepare_response(self, request, cached):
"""Verify our vary headers match and construct a real urllib3
HTTPResponse object.
"""
# Special case the '*' Vary value as it means we cannot actually
# determine if the cached response is suitable for this request.
# This case is also handled in the controller code when creating
# a cache entry, but is left here for backwards compatibility.
if "*" in cached.get("vary", {}):
return
# Ensure that the Vary headers for the cached response match our
# request
for header, value in cached.get("vary", {}).items():
if request.headers.get(header, None) != value:
return
body_raw = cached["response"].pop("body")
headers = CaseInsensitiveDict(data=cached["response"]["headers"])
if headers.get("transfer-encoding", "") == "chunked":
headers.pop("transfer-encoding")
cached["response"]["headers"] = headers
try:
body = io.BytesIO(body_raw)
except TypeError:
# This can happen if cachecontrol serialized to v1 format (pickle)
# using Python 2. A Python 2 str(byte string) will be unpickled as
# a Python 3 str (unicode string), which will cause the above to
# fail with:
#
# TypeError: 'str' does not support the buffer interface
body = io.BytesIO(body_raw.encode("utf8"))
return HTTPResponse(body=body,
preload_content=False,
**cached["response"])
def read_resp_record(self, resp_headers, payload):
len_ = payload.tell()
payload.seek(0)
warc_headers = self.parser.parse(payload)
warc_headers = CaseInsensitiveDict(warc_headers.headers)
record_type = warc_headers.get('WARC-Type', 'response')
if record_type == 'response':
status_headers = self.parser.parse(payload)
else:
status_headers = None
record = ArcWarcRecord('warc', record_type, warc_headers, payload,
status_headers, '', len_)
if record_type == 'response':
self._set_header_buff(record)
self.ensure_digest(record)
return record_type, record
def prepare_response(self, request, cached):
"""Verify our vary headers match and construct a real urllib3
HTTPResponse object.
"""
# Special case the '*' Vary value as it means we cannot actually
# determine if the cached response is suitable for this request.
# This case is also handled in the controller code when creating
# a cache entry, but is left here for backwards compatibility.
if "*" in cached.get("vary", {}):
return
# Ensure that the Vary headers for the cached response match our
# request
for header, value in cached.get("vary", {}).items():
if request.headers.get(header, None) != value:
return
body_raw = cached["response"].pop("body")
headers = CaseInsensitiveDict(data=cached["response"]["headers"])
if headers.get("transfer-encoding", "") == "chunked":
headers.pop("transfer-encoding")
cached["response"]["headers"] = headers
try:
body = io.BytesIO(body_raw)
except TypeError:
# This can happen if cachecontrol serialized to v1 format (pickle)
# using Python 2. A Python 2 str(byte string) will be unpickled as
# a Python 3 str (unicode string), which will cause the above to
# fail with:
#
# TypeError: 'str' does not support the buffer interface
body = io.BytesIO(body_raw.encode("utf8"))
return HTTPResponse(body=body, preload_content=False, **cached["response"])
def forward(self, method):
data = self.data_bytes
path = self.path
forward_headers = CaseInsensitiveDict(self.headers)
# force close connection
connection_header = forward_headers.get('Connection') or ''
if connection_header.lower() not in ['keep-alive', '']:
self.close_connection = 1
client_address = self.client_address[0]
server_address = ':'.join(map(str, self.server.server_address))
try:
# run the actual response forwarding
response = modify_and_forward(
method=method,
path=path,
data_bytes=data,
headers=forward_headers,
forward_base_url=self.proxy.forward_base_url,
listeners=self._listeners(),
request_handler=self,
client_address=client_address,
server_address=server_address)
# copy headers and return response
self.send_response(response.status_code)
# set content for chunked encoding
is_chunked = uses_chunked_encoding(response)
if is_chunked:
response._content = create_chunked_data(response._content)
# send headers
content_length_sent = False
for header_key, header_value in iteritems(response.headers):
# filter out certain headers that we don't want to transmit
if header_key.lower() not in ('transfer-encoding', 'date',
'server'):
self.send_header(header_key, header_value)
content_length_sent = content_length_sent or header_key.lower(
) == 'content-length'
# fix content-type header if needed
if not content_length_sent and not is_chunked:
self.send_header(
'Content-Length',
'%s' % len(response.content) if response.content else 0)
if isinstance(response, LambdaResponse):
self.send_multi_value_headers(response.multi_value_headers)
self.end_headers()
if response.content and len(response.content):
self.wfile.write(to_bytes(response.content))
except Exception as e:
trace = str(traceback.format_exc())
conn_errors = ('ConnectionRefusedError', 'NewConnectionError',
'Connection aborted', 'Unexpected EOF',
'Connection reset by peer',
'cannot read from timed out object')
conn_error = any(e in trace for e in conn_errors)
error_msg = 'Error forwarding request: %s %s' % (e, trace)
if 'Broken pipe' in trace:
LOG.warn(
'Connection prematurely closed by client (broken pipe).')
elif not self.proxy.quiet or not conn_error:
LOG.error(error_msg)
if is_local_test_mode():
# During a test run, we also want to print error messages, because
# log messages are delayed until the entire test run is over, and
# hence we are missing messages if the test hangs for some reason.
print('ERROR: %s' % error_msg)
self.send_response(502) # bad gateway
self.end_headers()
# force close connection
self.close_connection = 1
finally:
try:
self.wfile.flush()
except Exception as e:
LOG.warning('Unable to flush write file: %s' % e)
class Fossor(object):
'''
Engine for running automated investigations using three plugin types: variables, checks, reports. The engine:
- Collects variables
- Runs checks using applicable variables
- Reports noteworthy findings
'''
def __init__(self):
self.log = logging.getLogger(__name__)
self.plugin_parent_classes = [
fossor.plugin.Plugin, fossor.variables.variable.Variable,
fossor.checks.check.Check, fossor.reports.report.Report
]
# Variables
self.variables = CaseInsensitiveDict()
self.add_variable('timeout', 600) # Default timeout
# Imported Plugins
self.variable_plugins = set()
self.check_plugins = set()
self.report_plugins = set()
self.add_plugins(
) # Adds all plugins located within the fossor module recursively
def _import_submodules_by_module(self, module) -> set:
"""
Import all submodules from a module.
Returns a set of modules
"""
if '__path__' not in dir(module):
return set() # No other submodules to import, this is a file
result = set()
for loader, name, is_pkg in pkgutil.iter_modules(path=module.__path__):
sub_module = importlib.import_module(module.__name__ + '.' + name)
result.add(sub_module)
result.update(self._import_submodules_by_module(sub_module))
return result
def _import_submodules_by_path(self, path: str) -> set:
"""
Import all submodules from a path.
Returns a set of modules
"""
# Get a list of all the python files
# Loop over list and group modules
path = os.path.normpath(path)
all_file_paths = []
for root, dirs, files in os.walk(path):
for file in files:
all_file_paths.append(os.path.join(root, file))
python_file_paths = [
file for file in all_file_paths
if '.py' == os.path.splitext(file)[1]
]
result = set()
for filepath in python_file_paths:
# Create a module_name
relative_path = filepath.split(path)[-1]
module_name = os.path.splitext(relative_path)[
0] # Remove file extension
module_name = module_name.lstrip(
os.path.sep) # Removing leading slash
module_name = module_name.replace(os.path.sep,
'.') # Change / to .
module_name = f'fossor.local.{module_name}'
# Import with the new name
spec = importlib.util.spec_from_file_location(
module_name, filepath)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
result.add(module)
return result
def add_plugins(self, source=fossor):
'''
Recursively return a dict of plugins (classes) that inherit from the given parent_class) from within that source.
source accepts either a python module or a filesystem path as a string.
'''
if source is None:
return
if type(source) == str:
modules = self._import_submodules_by_path(path=source)
else:
modules = self._import_submodules_by_module(source)
for module in modules:
for obj_name, obj in module.__dict__.items():
# Get objects from each module that look like plugins for the Check abstract class
if not inspect.isclass(
obj
): # Must check if a class otherwise the next check with issubclass will get a TypeError
continue
if obj in self.plugin_parent_classes: # If this is one of the parent classes, skip, we only want the children
continue
if issubclass(obj, fossor.variables.variable.Variable):
self.variable_plugins.add(obj)
elif issubclass(obj, fossor.checks.check.Check):
self.check_plugins.add(obj)
elif issubclass(obj, fossor.reports.report.Report):
self.report_plugins.add(obj)
else:
continue
self.log.debug(f"Adding Plugin ({obj_name}, {obj})")
def clear_plugins(self):
self.variable_plugins = set()
self.check_plugins = set()
self.report_plugins = set()
def list_plugins(self):
variables = list(self.variable_plugins)
checks = list(self.check_plugins)
reports = list(self.report_plugins)
return variables + checks + reports
def _terminate_process_group(self, process):
'''Kill a process group leader and its process group'''
pid = process.pid
pgid = os.getpgid(pid)
if pid != pgid:
raise ValueError(
f"Trying to terminate a pid that is not a pgid leader. pid: {pid}, pgid: {pgid}."
) # This should never happen
# Try these signals to kill the process group
for s in (signal.SIGTERM, signal.SIGKILL):
try:
children = psutil.Process().children(recursive=True)
except psutil_exceptions:
continue # Process is already dead
if len(children) > 0:
try:
self.log.warning(
f"Issuing signal {s} to process group {pgid} because it has run for too long."
)
os.killpg(pgid, s)
time.sleep(1)
except psutil_exceptions:
pass
except PermissionError: # Occurs if process has already exited on MacOS
pass
# Terminate the plugin process (which is likely dead at this point)
process.terminate()
process.join()
def _run_plugins_parallel(self, plugins):
'''
Accepts function to run, and a dictionary of plugins to run. Format is name:module
Returns an output queue which outputs two values: plugin-name, output. Queue ends when plugin-name is EOF.
'''
def _run_plugins_parallel_helper(timeout):
setproctitle.setproctitle('Plugin-Runner')
processes = []
queue_lock = mp.Lock()
for Plugin in plugins:
process = mp.Process(target=self.run_plugin,
name=Plugin.get_name(),
args=(Plugin, output_queue, queue_lock))
process.start()
processes.append(process)
start_time = time.time()
should_end = False
while len(processes) > 0:
time_spent = time.time() - start_time
should_end = time_spent > timeout
try:
os.kill(os.getppid(),
0) # Check if the parent process is alive
except ProcessLookupError:
should_end = True
for process in processes: # Using list so I can delete while iterating
process_is_dead = not process.is_alive(
) and process.exitcode is not None
if process_is_dead:
process.join()
processes.remove(process)
continue
if should_end:
with queue_lock:
self.log.error(
f"Process {process} for plugin {process.name} ran longer than the timeout period: {timeout}"
)
if self.variables.get('verbose'):
output_queue.put((
process.name,
'Timed out (use --time-out to increase timeout)'
))
self._terminate_process_group(process)
processes.remove(process)
time.sleep(.1)
with queue_lock:
output_queue.put(('Stats', f'Ran {len(plugins)} plugins.'))
output_queue.put(
('EOF', 'EOF')) # Indicate the queue is finished
return output_queue
timeout = int(self.variables.get('timeout'))
# Run plugins in parallel
# This child process will spawn child processes for each plugin, and then do the final termination and clean-up.
# A separate child process for spawning plugin processes is needed because only a parent process can terminate its children.
# This allows us to continue to the reporting phase while the checks finish.
output_queue = mp.Queue()
parallel_plugins_helper_process = mp.Process(
target=_run_plugins_parallel_helper,
name='parallel_plugins_helper_process',
args=(timeout, ))
parallel_plugins_helper_process.start()
return output_queue
def run_plugin(self, Plugin, output_queue, queue_lock):
try:
os.setsid(
) # Causes this plugin to become a process group leader, we can then kill that process group to kill all potential subprocesses
p = Plugin()
setproctitle.setproctitle(p.get_name())
output = ''
if p.should_run():
output = p.run_helper(variables=self.variables)
if output:
with queue_lock: # Prevents the queue from becoming corrupt on process termination
output_queue.put((p.get_name(), output))
except Exception as e:
# This ensures that any plugins that fail to initialize show up properly in the report and don't output errors mid report.
self.log.exception(f"Plugin {p} failed to initialize", e)
def get_variables(self):
'''
Gather all possible variables that may be useful for plugins
- Variables can have interdependencies, this will continue gathering missing variables
until no new variables have appeared after a run
- Can be run again if new variables are added to fill in any blanks
'''
while True:
variable_count = len(self.variables)
output_queue = self._run_plugins_parallel(self.variable_plugins)
while True:
name, output = output_queue.get()
if 'EOF' in name:
break
if name in self.variables:
self.log.debug(
"Skipping variable plugin because it has already been found: {vp}"
.format(vp=name))
continue
self.add_variable(name, output)
output_queue.close()
if variable_count == len(self.variables):
self.log.debug("Done finding variables: {variables}".format(
variables=self.variables))
break
def _convert_simple_type(self, value):
'''Convert String to simple type if possible'''
if type(value) == str:
if 'false' == value.lower():
return False
if 'true' == value.lower():
return True
try:
return int(value)
except ValueError:
pass
try:
return float(value)
except ValueError:
pass
return value
def add_variables(self, **kwargs):
for name, value in kwargs.items():
self.add_variable(name=name, value=value)
return True
def add_variable(self, name, value):
'''Adds variable, converts numbers and booleans to their types'''
if name in self.variables:
old_value = self.variables[name]
if value != old_value:
self.log.warning(
f"Variable {name} with value of {old_value} is being replaced with {value}"
)
value = self._convert_simple_type(value)
self.variables[name] = value
self.log.debug(
"Added variable {name} with value of {value} (type={type})".format(
name=name, value=value, type=type(value)))
return True
def _process_whitelist(self, whitelist):
if whitelist:
whitelist = [name.casefold() for name in whitelist]
self.variable_plugins = {
plugin
for plugin in self.variable_plugins
if plugin.get_name().casefold() in whitelist
}
self.check_plugins = {
plugin
for plugin in self.check_plugins
if plugin.get_name().casefold() in whitelist
}
def _process_blacklist(self, blacklist):
if blacklist:
blacklist = [name.casefold() for name in blacklist]
self.variable_plugins = {
plugin
for plugin in self.variable_plugins
if plugin.get_name().casefold() not in blacklist
}
self.check_plugins = {
plugin
for plugin in self.check_plugins
if plugin.get_name().casefold() not in blacklist
}
def run(self, report='StdOut', **kwargs):
'''Runs Fossor with the given report. Method returns a string of the report output.'''
self.log.debug("Starting Fossor")
self._process_whitelist(kwargs.get('whitelist'))
self._process_blacklist(kwargs.get('blacklist'))
# Add kwargs as variables
self.add_variables(**kwargs)
# Add directory plugins
self.add_plugins(kwargs.get('plugin_dir'))
# Gather Variables
self.get_variables()
# Run checks
output_queue = self._run_plugins_parallel(self.check_plugins)
# Run Report
try:
Report_Plugin = [
plugin for plugin in self.report_plugins
if report.lower() == plugin.get_name().lower()
].pop()
except IndexError:
message = f"Report_Plugin: {report}, was not found. Possibly reports are: {[plugin.get_name() for plugin in self.report_plugins]}"
self.log.critical(message)
raise ValueError(message)
report_plugin = Report_Plugin()
return report_plugin.run(variables=self.variables,
report_input=output_queue)
def forward(self, method):
data = self.data_bytes
forward_headers = CaseInsensitiveDict(self.headers)
# force close connection
if forward_headers.get('Connection', '').lower() != 'keep-alive':
self.close_connection = 1
path = self.path
if '://' in path:
path = '/' + path.split('://', 1)[1].split('/', 1)[1]
forward_url = self.proxy.forward_url
for listener in self._listeners():
if listener:
forward_url = listener.get_forward_url(
method, path, data, forward_headers) or forward_url
proxy_url = '%s%s' % (forward_url, path)
target_url = self.path
if '://' not in target_url:
target_url = '%s%s' % (forward_url, target_url)
# update original "Host" header (moto s3 relies on this behavior)
if not forward_headers.get('Host'):
forward_headers['host'] = urlparse(target_url).netloc
if 'localhost.atlassian.io' in forward_headers.get('Host'):
forward_headers['host'] = 'localhost'
forward_headers['X-Forwarded-For'] = self.build_x_forwarded_for(
forward_headers)
try:
response = None
modified_request = None
# update listener (pre-invocation)
for listener in self._listeners():
if not listener:
continue
listener_result = listener.forward_request(
method=method,
path=path,
data=data,
headers=forward_headers)
if isinstance(listener_result, Response):
response = listener_result
break
if isinstance(listener_result, dict):
response = Response()
response._content = json.dumps(listener_result)
response.status_code = 200
break
elif isinstance(listener_result, Request):
modified_request = listener_result
data = modified_request.data
forward_headers = modified_request.headers
break
elif listener_result is not True:
# get status code from response, or use Bad Gateway status code
code = listener_result if isinstance(listener_result,
int) else 503
self.send_response(code)
self.send_header('Content-Length', '0')
# allow pre-flight CORS headers by default
self._send_cors_headers()
self.end_headers()
return
# perform the actual invocation of the backend service
if response is None:
forward_headers['Connection'] = forward_headers.get(
'Connection') or 'close'
data_to_send = self.data_bytes
request_url = proxy_url
if modified_request:
if modified_request.url:
request_url = '%s%s' % (forward_url,
modified_request.url)
data_to_send = modified_request.data
response = self.method(request_url,
data=data_to_send,
headers=forward_headers,
stream=True)
# prevent requests from processing response body
if not response._content_consumed and response.raw:
response._content = response.raw.read()
# update listener (post-invocation)
if self.proxy.update_listener:
kwargs = {
'method': method,
'path': path,
'data': data,
'headers': forward_headers,
'response': response
}
if 'request_handler' in inspect.getargspec(
self.proxy.update_listener.return_response)[0]:
# some listeners (e.g., sqs_listener.py) require additional details like the original
# request port, hence we pass in a reference to this request handler as well.
kwargs['request_handler'] = self
updated_response = self.proxy.update_listener.return_response(
**kwargs)
if isinstance(updated_response, Response):
response = updated_response
# copy headers and return response
self.send_response(response.status_code)
content_length_sent = False
for header_key, header_value in iteritems(response.headers):
# filter out certain headers that we don't want to transmit
if header_key.lower() not in ('transfer-encoding', 'date',
'server'):
self.send_header(header_key, header_value)
content_length_sent = content_length_sent or header_key.lower(
) == 'content-length'
if not content_length_sent:
self.send_header(
'Content-Length',
'%s' % len(response.content) if response.content else 0)
# allow pre-flight CORS headers by default
self._send_cors_headers(response)
self.end_headers()
if response.content and len(response.content):
self.wfile.write(to_bytes(response.content))
except Exception as e:
trace = str(traceback.format_exc())
conn_errors = ('ConnectionRefusedError', 'NewConnectionError',
'Connection aborted', 'Unexpected EOF',
'Connection reset by peer')
conn_error = any(e in trace for e in conn_errors)
error_msg = 'Error forwarding request: %s %s' % (e, trace)
if 'Broken pipe' in trace:
LOG.warn(
'Connection prematurely closed by client (broken pipe).')
elif not self.proxy.quiet or not conn_error:
LOG.error(error_msg)
if os.environ.get(ENV_INTERNAL_TEST_RUN):
# During a test run, we also want to print error messages, because
# log messages are delayed until the entire test run is over, and
# hence we are missing messages if the test hangs for some reason.
print('ERROR: %s' % error_msg)
self.send_response(502) # bad gateway
self.end_headers()
# force close connection
self.close_connection = 1
finally:
try:
self.wfile.flush()
except Exception as e:
LOG.warning('Unable to flush write file: %s' % e)
def post(self):
form = RequestDataForm()
data = form.data.data
lines = data.splitlines(True)
if len(lines) < 3:
return 'data less 3 lines'
origin_headers = []
body = []
body_start = False
for index, line in enumerate(lines):
if index == 0:
method, path, _ = line.split(' ')
continue
if not line.split():
body_start = True
continue
if body_start:
body.append(line)
else:
line = line.strip()
key, value = line.split(': ', 1)
origin_headers.append((key, value))
# for get header value
header_dict = CaseInsensitiveDict(origin_headers)
method = method.lower()
body = ''.join(body)
content_type = header_dict.get('Content-Type', '')
# set headers
headers = []
origin_host = header_dict.get('Host')
if form.host.data and origin_host and form.host.data != origin_host:
headers.append(('Host', header_dict.get('Host')))
user_agent = header_dict.get('User-Agent')
referer = header_dict.get('Referer')
if user_agent:
headers.append(('User-Agent', user_agent))
if referer:
headers.append(('Referer', referer))
# set cookie
cookies = []
cookie = header_dict.get('Cookie')
C = SimpleCookie(cookie)
for morsel in C.values():
cookies.append((morsel.key, morsel.coded_value))
host = form.host.data or header_dict.get('Host')
p = urlsplit(path)
url = urljoin('http://{}'.format(host), p.path)
params = [(x, repr_value(y)) for x, y in parse_qsl(p.query)]
if not content_type:
pass
elif 'x-www-form-urlencoded' in content_type:
body = [(x, repr_value(y)) for x, y in parse_qsl(body)]
elif 'json' in content_type:
body = [(x, repr_value(y)) for x, y in json.loads(body).items()]
else:
headers.append(('Content-Type', content_type))
code = render_template(
'code.html',
method=method,
url=url,
params=params,
body=body,
headers=headers,
cookies=cookies,
content_type=content_type
)
return render_template('code-page.html', code=code)
def match(self, request):
"""Does this URI match the request?"""
# Match HTTP verb - GET, POST, PUT, DELETE, etc.
if self.method is not None:
if request.command.lower() != self.method.lower():
return False
# Match path
if self.path != urlparse.urlparse(request.path).path:
return False
# Match querystring
if request.querystring() != self.querystring:
return False
# Match headers
if self.headers is not None:
for header_var, header_value in self.headers.items():
request_headers = CaseInsensitiveDict(request.headers)
if request_headers.get(header_var) is None:
return False
req_maintext, req_pdict = cgi.parse_header(request_headers.get(header_var))
mock_maintext, mock_pdict = cgi.parse_header(header_value)
if "boundary" in req_pdict and "boundary" in mock_pdict:
req_pdict['boundary'] = "xxx"
mock_pdict['boundary'] = "xxx"
if req_maintext != mock_maintext:
return False
if mock_pdict != {}:
if req_pdict != mock_pdict:
return False
# Match processed request data
if self.request_data is not None:
# Check if exact match before parsing
if request.body != self.request_data:
if self.request_content_type.startswith("application/json"):
if json.loads(request.body) != json.loads(self.request_data):
return False
elif self.request_content_type.startswith("application/x-www-form-urlencoded"):
if parse_qsl_as_dict(request.body) != parse_qsl_as_dict(self.request_data):
return False
elif self.request_content_type.startswith('application/xml'):
actual_request_data_root = etree.fromstring(request.body)
mock_request_data_root = etree.fromstring(self.request_data)
if not xml_elements_equal(actual_request_data_root, mock_request_data_root):
return False
elif self.request_content_type.startswith("multipart/form-data"):
ctype, pdict = cgi.parse_header(request.headers.get('Content-Type'))
req_multipart = cgi.parse_multipart(
io.BytesIO(request.body.encode('utf8')),
{x: y.encode('utf8') for x, y in pdict.items()}
)
ctype, pdict = cgi.parse_header(self.headers.get('Content-Type'))
mock_multipart = cgi.parse_multipart(
io.BytesIO(self.request_data.encode('utf8')),
{x: y.encode('utf8') for x, y in pdict.items()}
)
return mock_multipart == req_multipart
else:
if request.body != self.request_data:
return False
return True
def forward(self, method):
path = self.path
if '://' in path:
path = '/' + path.split('://', 1)[1].split('/', 1)[1]
proxy_url = '%s%s' % (self.proxy.forward_url, path)
target_url = self.path
if '://' not in target_url:
target_url = '%s%s' % (self.proxy.forward_url, target_url)
data = self.data_bytes
forward_headers = CaseInsensitiveDict(self.headers)
# update original "Host" header (moto s3 relies on this behavior)
if not forward_headers.get('Host'):
forward_headers['host'] = urlparse(target_url).netloc
if 'localhost.atlassian.io' in forward_headers.get('Host'):
forward_headers['host'] = 'localhost'
try:
response = None
modified_request = None
# update listener (pre-invocation)
if self.proxy.update_listener:
listener_result = self.proxy.update_listener.forward_request(method=method,
path=path, data=data, headers=forward_headers)
if isinstance(listener_result, Response):
response = listener_result
elif isinstance(listener_result, Request):
modified_request = listener_result
data = modified_request.data
forward_headers = modified_request.headers
elif listener_result is not True:
# get status code from response, or use Bad Gateway status code
code = listener_result if isinstance(listener_result, int) else 503
self.send_response(code)
self.end_headers()
return
# perform the actual invocation of the backend service
if response is None:
if modified_request:
response = self.method(proxy_url, data=modified_request.data,
headers=modified_request.headers)
else:
response = self.method(proxy_url, data=self.data_bytes,
headers=forward_headers)
# update listener (post-invocation)
if self.proxy.update_listener:
kwargs = {
'method': method,
'path': path,
'data': data,
'headers': forward_headers,
'response': response
}
if 'request_handler' in inspect.getargspec(self.proxy.update_listener.return_response)[0]:
# some listeners (e.g., sqs_listener.py) require additional details like the original
# request port, hence we pass in a reference to this request handler as well.
kwargs['request_handler'] = self
updated_response = self.proxy.update_listener.return_response(**kwargs)
if isinstance(updated_response, Response):
response = updated_response
# copy headers and return response
self.send_response(response.status_code)
content_length_sent = False
for header_key, header_value in iteritems(response.headers):
# filter out certain headers that we don't want to transmit
if header_key.lower() not in ('transfer-encoding', 'date', 'server'):
self.send_header(header_key, header_value)
content_length_sent = content_length_sent or header_key.lower() == 'content-length'
if not content_length_sent:
self.send_header('Content-Length', '%s' % len(response.content) if response.content else 0)
# allow pre-flight CORS headers by default
if 'Access-Control-Allow-Origin' not in response.headers:
self.send_header('Access-Control-Allow-Origin', '*')
if 'Access-Control-Allow-Methods' not in response.headers:
self.send_header('Access-Control-Allow-Methods', ','.join(CORS_ALLOWED_METHODS))
if 'Access-Control-Allow-Headers' not in response.headers:
self.send_header('Access-Control-Allow-Headers', ','.join(CORS_ALLOWED_HEADERS))
self.end_headers()
if response.content and len(response.content):
self.wfile.write(to_bytes(response.content))
self.wfile.flush()
except Exception as e:
trace = str(traceback.format_exc())
conn_errors = ('ConnectionRefusedError', 'NewConnectionError')
conn_error = any(e in trace for e in conn_errors)
error_msg = 'Error forwarding request: %s %s' % (e, trace)
if 'Broken pipe' in trace:
LOGGER.warn('Connection prematurely closed by client (broken pipe).')
elif not self.proxy.quiet or not conn_error:
LOGGER.error(error_msg)
if os.environ.get(ENV_INTERNAL_TEST_RUN):
# During a test run, we also want to print error messages, because
# log messages are delayed until the entire test run is over, and
# hence we are missing messages if the test hangs for some reason.
print('ERROR: %s' % error_msg)
self.send_response(502) # bad gateway
self.end_headers()
def proxy(identifier, in_method, in_headers, data):
# find endpoint
endpoint = pg.select1(
'endpoints',
what=['definition', 'method', 'pass_headers',
'headers', 'url', 'url_dynamic'],
where={'id': identifier}
)
if not endpoint:
return 'endpoint not found, create it at ' \
'<a href="/dashboard">dashboard</a>', 404, {}
event = {
'in': {
'time': time.time(),
'method': in_method,
'headers': dict(in_headers),
'body': data,
'replay': True
},
'out': {
'method': endpoint['method'],
'url': endpoint['url'],
'body': None,
'headers': {}
},
'response': {
'code': 0,
'body': ''
}
}
mutated, error = jq(endpoint['definition'], data=data)
if not mutated or error:
event['out']['error'] = error.decode('utf-8')
publish(identifier, event)
return 'transmutated into null and aborted', 201, {}
h = CaseInsensitiveDict({'Content-Type': 'application/json'})
if endpoint['pass_headers']:
h.update(in_headers)
h.update(endpoint['headers'])
event['out']['headers'] = dict(h)
# reformat the mutated data
mutatedjson = json.loads(mutated.decode('utf-8'))
if h.get('content-type') == 'application/x-www-form-urlencoded':
# oops, not json
mutated = urlencode(mutatedjson)
else:
mutated = json.dumps(mutatedjson)
event['out']['body'] = mutated
if endpoint['url_dynamic']:
urlb, error = jq(endpoint['url'], data=data)
print('URL', urlb, 'ERROR', error)
if not urlb:
event['out']['url_error'] = error.decode('utf-8')
publish(identifier, event)
return 'url building has failed', 200, {}
url = urlb.decode('utf-8')
event['out']['url'] = url
else:
url = endpoint['url']
# event['out'] is completed at this point
# and we all have everything needed to perform the request
if url and is_valid_url(url):
# we will only perform a request if there is an URL and it is valid
try:
s = requests.Session()
req = requests.Request(endpoint['method'], url,
data=mutated, headers=h).prepare()
resp = s.send(req, timeout=4)
if not resp.ok:
print('FAILED TO POST', resp.text, identifier, mutated)
except requests.exceptions.RequestException as e:
print(identifier, 'FAILED TO POST', mutated, 'TO URL', url)
print(e)
publish(identifier, event)
return "<request failed: '%s'>" % e, 503, {}
event['response']['code'] = resp.status_code
event['response']['body'] = resp.text
publish(identifier, event)
return resp.text, resp.status_code, dict(resp.headers)
else:
# no valid URL, just testing
publish(identifier, event)
return 'no URL to send this to', 201, {}
def cache_response(self, request, response, body=None, status_codes=None):
"""
Algorithm for caching requests.
This assumes a requests Response object.
"""
# From httplib2: Don't cache 206's since we aren't going to
# handle byte range requests
cacheable_status_codes = status_codes or self.cacheable_status_codes
if response.status not in cacheable_status_codes:
logger.debug(
"Status code %s not in %s", response.status, cacheable_status_codes
)
return
response_headers = CaseInsensitiveDict(response.headers)
# If we've been given a body, our response has a Content-Length, that
# Content-Length is valid then we can check to see if the body we've
# been given matches the expected size, and if it doesn't we'll just
# skip trying to cache it.
if (
body is not None
and "content-length" in response_headers
and response_headers["content-length"].isdigit()
and int(response_headers["content-length"]) != len(body)
):
return
cc_req = self.parse_cache_control(request.headers)
cc = self.parse_cache_control(response_headers)
cache_url = self.cache_url(request.url)
logger.debug('Updating cache with response from "%s"', cache_url)
# Delete it from the cache if we happen to have it stored there
no_store = False
if "no-store" in cc:
no_store = True
logger.debug('Response header has "no-store"')
if "no-store" in cc_req:
no_store = True
logger.debug('Request header has "no-store"')
if no_store and self.cache.get(cache_url):
logger.debug('Purging existing cache entry to honor "no-store"')
self.cache.delete(cache_url)
if no_store:
return
# https://tools.ietf.org/html/rfc7234#section-4.1:
# A Vary header field-value of "*" always fails to match.
# Storing such a response leads to a deserialization warning
# during cache lookup and is not allowed to ever be served,
# so storing it can be avoided.
if "*" in response_headers.get("vary", ""):
logger.debug('Response header has "Vary: *"')
return
# If we've been given an etag, then keep the response
if self.cache_etags and "etag" in response_headers:
logger.debug("Caching due to etag")
self.cache.set(
cache_url, self.serializer.dumps(request, response, body=body)
)
# Add to the cache any 301s. We do this before looking that
# the Date headers.
elif response.status == 301:
logger.debug("Caching permanant redirect")
self.cache.set(cache_url, self.serializer.dumps(request, response))
# Add to the cache if the response headers demand it. If there
# is no date header then we can't do anything about expiring
# the cache.
elif "date" in response_headers:
# cache when there is a max-age > 0
if "max-age" in cc and cc["max-age"] > 0:
logger.debug("Caching b/c date exists and max-age > 0")
self.cache.set(
cache_url, self.serializer.dumps(request, response, body=body)
)
# If the request can expire, it means we should cache it
# in the meantime.
elif "expires" in response_headers:
if response_headers["expires"]:
logger.debug("Caching b/c of expires header")
self.cache.set(
cache_url, self.serializer.dumps(request, response, body=body)
)
class MessageParser(object):
def __init__(self, name, content, init_headers={}):
self.name = name
self.content = content
self.headers = CaseInsensitiveDict(init_headers)
self.html_stats = HTMLStats()
self.email_stats = EmailStats()
if '\n\n' in self.content:
self.head, self.body = self.content.split('\n\n', 1)
else:
self.head, self.body = "", ""
self._parse_headers()
self._decode_body()
self._count_links()
@property
def subject(self):
return self.headers['Subject']
@property
def mime_type(self):
if 'Content-Type' in self.headers:
content_type = self.headers['Content-Type']
mime_type = content_type.split(';', 1)[0]
mime_type = mime_type.strip().lower()
return mime_type
@property
def charset(self):
if 'Content-Type' in self.headers:
content_type = self.headers['Content-Type']
rem = RE_CHARSET.search(content_type)
if rem:
charset = rem.group(1).lower()
return charset
@property
def is_multipart(self):
if self.mime_type and self.mime_type.startswith('multipart/'):
return True
else:
return False
def _get_boundary(self):
if 'Content-Type' in self.headers:
content_type = self.headers['Content-Type']
rem = RE_HEADER_BOUNDARY.search(content_type)
if rem:
return rem.group(1)
def _parse_headers(self):
lines = self.head.split('\n')
# Drop the "From..." line
while lines and not RE_HEADER.match(lines[0]):
del lines[0]
while lines:
line = self._get_header_line(lines)
key, value = line.split(':', 1)
key, value = key.strip(), value.strip()
self.headers[key] = value
def _get_header_line(self, lines):
line_parts = [lines.pop(0)]
while lines and (lines[0].startswith('\t') or
lines[0].startswith(' ')):
line_parts.append(lines.pop(0))
line = " ".join([part.strip() for part in line_parts])
return line
def _decode_body(self):
if self.mime_type and (self.mime_type.startswith('image/') or
self.mime_type.startswith('application/')):
LOG.info("Body marked as image, skipping body")
self.email_stats['attached_images'] += 1
self.body = ""
return
if self.is_multipart:
LOG.debug("Detected multipart/* content-type")
self._decode_multipart_body()
else:
self._decode_single_body()
def _decode_single_body(self):
self.body = self.body.strip()
cte = self.headers.get('Content-Transfer-Encoding', '').lower()
if 'quoted-printable' in cte:
LOG.debug("Detected quoted-printable encoding, decoding")
self.body = quopri.decodestring(self.body)
if 'base64' in cte:
LOG.debug("Detected base64 encoding, decoding")
try:
self.body = base64.decodestring(self.body)
except base64.binascii.Error:
LOG.info("base64 decoder failed, trying partial decoding")
self.body = base64_partial_decode(self.body)
LOG.debug("Detected charset: %s", self.charset)
try:
self.body = self.body.decode(
validate_charset(self.charset) and self.charset or 'ascii',
'strict'
)
except UnicodeDecodeError:
LOG.info('Error during strict decoding')
self.email_stats['charset_errors'] = 1
self.body = self.body.decode(
validate_charset(self.charset) and self.charset or 'ascii',
'ignore'
)
if self._guess_html():
LOG.debug("Message recognized as HTML")
self._parse_html()
else:
LOG.debug("Message recognized as plaintext")
def _decode_multipart_body(self):
boundary = self._get_boundary()
if not boundary:
LOG.warn("Message detected as multipart but boundary "
"declaration was not found")
return
start_bnd = '\n' + '--' + boundary + '\n'
end_bnd = '\n' + '--' + boundary + '--' + '\n'
self.body = '\n' + self.body # for string matching purpose
try:
start_index = self.body.index(start_bnd) + len(start_bnd)
except ValueError:
LOG.warn("Cannot find boundaries in body, "
"treating as single message")
self._decode_single_body()
return
end_index = self.body.rfind(end_bnd)
if end_index < 0:
end_index = None
content = self.body[start_index:end_index]
parts = content.split(start_bnd)
messages = [MessageParser(self.name, msg_content, self.headers)
for msg_content in parts]
self.body = "\n".join([msg.body for msg in messages])
for msg in messages:
self.email_stats += msg.email_stats
self.html_stats += msg.html_stats
def _guess_html(self):
if self.mime_type and self.mime_type == 'text/html':
return True
else:
return False
def _parse_html(self):
parser = StatsHTMLParser()
parser.feed(self.body)
parser.close()
self.body = parser.text
self.html_stats = parser.stats
def _count_links(self):
mail_links = RE_MAIL_LINK.findall(self.body)
self.body = RE_MAIL_LINK.sub('', self.body)
self.email_stats['mail_links'] = len(mail_links)
http_links = RE_HTTP_LINK.findall(self.body)
self.body = RE_HTTP_LINK.sub('', self.body)
self.email_stats['http_links'] = len(http_links)
http_raw_links = [link for link in http_links
if RE_HTTP_RAW_LINK.match(link)]
self.email_stats['http_links'] = len(http_raw_links)
def as_series(self):
body_length = len(self.body)
s = pd.Series({'plain_body': self.body, 'body_length': body_length})
rel_email_stats = self.email_stats.astype('float32')
rel_email_stats.index = ['rel_' + idx for idx in
rel_email_stats.index]
if body_length > 0:
rel_email_stats /= body_length
else:
rel_email_stats[:] = 0
rel_html_stats = self.html_stats.astype('float32')
rel_html_stats.index = ['rel_' + idx for idx in
rel_html_stats.index]
if body_length > 0:
rel_html_stats /= body_length
else:
rel_html_stats[:] = 0
s_headers = pd.Series(dict(self.headers))
s_headers.index = s_headers.index.map(lambda h: 'h_' + h.lower())
s = (s
.append(s_headers, True)
.append(self.email_stats, True)
.append(self.html_stats, True)
.append(rel_email_stats, True)
.append(rel_html_stats, True)
)
s.name = self.name
s = s.fillna(0)
return s
class MockRestURI(object):
"""Representation of a mock URI."""
def __init__(self, uri_dict):
self._uri_dict = uri_dict
self.name = uri_dict.get('name', None)
self.fullpath = uri_dict['request']['path']
self._regexp = False
self.path = urlparse.urlparse(self.fullpath).path
self.querystring = urlparse.parse_qs(
urlparse.urlparse(self.fullpath).query,
keep_blank_values=True
)
self.method = uri_dict['request'].get('method', None)
self.headers = CaseInsensitiveDict(uri_dict['request'].get('headers', {}))
self.return_code = int(uri_dict['response'].get('code', '200'))
self.request_content_type = self.headers.get('Content-Type', '')
self.response_location = uri_dict['response'].get('location', None)
self.response_content = uri_dict['response'].get('content', "")
self.wait = float(uri_dict['response'].get('wait', 0.0))
self.request_data = uri_dict['request'].get('data', None)
self.encoding = uri_dict['request'].get("encoding", None)
self.response_headers = uri_dict['response'].get("headers", {})
def match(self, request):
"""Does this URI match the request?"""
# Match HTTP verb - GET, POST, PUT, DELETE, etc.
if self.method is not None:
if request.command.lower() != self.method.lower():
return False
# Match path
if self.path != urlparse.urlparse(request.path).path:
return False
# Match querystring
if request.querystring() != self.querystring:
return False
# Match headers
if self.headers is not None:
for header_var, header_value in self.headers.items():
request_headers = CaseInsensitiveDict(request.headers)
if request_headers.get(header_var) is None:
return False
req_maintext, req_pdict = cgi.parse_header(request_headers.get(header_var))
mock_maintext, mock_pdict = cgi.parse_header(header_value)
if "boundary" in req_pdict and "boundary" in mock_pdict:
req_pdict['boundary'] = "xxx"
mock_pdict['boundary'] = "xxx"
if req_maintext != mock_maintext:
return False
if mock_pdict != {}:
if req_pdict != mock_pdict:
return False
# Match processed request data
if self.request_data is not None:
# Check if exact match before parsing
if request.body != self.request_data:
if self.request_content_type.startswith("application/json"):
if json.loads(request.body) != json.loads(self.request_data):
return False
elif self.request_content_type.startswith("application/x-www-form-urlencoded"):
if parse_qsl_as_dict(request.body) != parse_qsl_as_dict(self.request_data):
return False
elif self.request_content_type.startswith('application/xml'):
actual_request_data_root = etree.fromstring(request.body)
mock_request_data_root = etree.fromstring(self.request_data)
if not xml_elements_equal(actual_request_data_root, mock_request_data_root):
return False
elif self.request_content_type.startswith("multipart/form-data"):
ctype, pdict = cgi.parse_header(request.headers.get('Content-Type'))
req_multipart = cgi.parse_multipart(
io.BytesIO(request.body.encode('utf8')),
{x: y.encode('utf8') for x, y in pdict.items()}
)
ctype, pdict = cgi.parse_header(self.headers.get('Content-Type'))
mock_multipart = cgi.parse_multipart(
io.BytesIO(self.request_data.encode('utf8')),
{x: y.encode('utf8') for x, y in pdict.items()}
)
return mock_multipart == req_multipart
else:
if request.body != self.request_data:
return False
return True
def querystring_string(self):
# TODO : Refactor this and docstring.
query = ''
for key in self.querystring.keys():
for item in self.querystring[key]:
query += str(key) + '=' + item + "&"
query = query.rstrip("&")
return "?" + query if query else ""
def example_path(self):
return xeger.xeger(self.path) if self._regexp else self.path + self.querystring_string()
def return_code_description(self):
return status_codes.CODES.get(self.return_code)[0]
def request_data_values(self):
if self.request_data is not None:
return self.request_data.get('values', {}).iteritems()
else:
return []
def request_data_type(self):
if self.request_data is not None:
return self.request_data.get('encoding')
else:
return None
class tizgmusicproxy(object):
"""A class for logging into a Google Play Music account and retrieving song
URLs.
"""
all_songs_album_title = "All Songs"
thumbs_up_playlist_name = "Thumbs Up"
# pylint: disable=too-many-instance-attributes,too-many-public-methods
def __init__(self, email, password, device_id):
self.__gmusic = Mobileclient()
self.__email = email
self.__device_id = device_id
self.logged_in = False
self.queue = list()
self.queue_index = -1
self.play_queue_order = list()
self.play_modes = TizEnumeration(["NORMAL", "SHUFFLE"])
self.current_play_mode = self.play_modes.NORMAL
self.now_playing_song = None
userdir = os.path.expanduser('~')
tizconfig = os.path.join(userdir, ".config/tizonia/." + email + ".auth_token")
auth_token = ""
if os.path.isfile(tizconfig):
with open(tizconfig, "r") as f:
auth_token = pickle.load(f)
if auth_token:
# 'Keep track of the auth token' workaround. See:
# https://github.com/diraimondo/gmusicproxy/issues/34#issuecomment-147359198
print_msg("[Google Play Music] [Authenticating] : " \
"'with cached auth token'")
self.__gmusic.android_id = device_id
self.__gmusic.session._authtoken = auth_token
self.__gmusic.session.is_authenticated = True
try:
self.__gmusic.get_registered_devices()
except CallFailure:
# The token has expired. Reset the client object
print_wrn("[Google Play Music] [Authenticating] : " \
"'auth token expired'")
self.__gmusic = Mobileclient()
auth_token = ""
if not auth_token:
attempts = 0
print_nfo("[Google Play Music] [Authenticating] : " \
"'with user credentials'")
while not self.logged_in and attempts < 3:
self.logged_in = self.__gmusic.login(email, password, device_id)
attempts += 1
with open(tizconfig, "a+") as f:
f.truncate()
pickle.dump(self.__gmusic.session._authtoken, f)
self.library = CaseInsensitiveDict()
self.song_map = CaseInsensitiveDict()
self.playlists = CaseInsensitiveDict()
self.stations = CaseInsensitiveDict()
def logout(self):
""" Reset the session to an unauthenticated, default state.
"""
self.__gmusic.logout()
def set_play_mode(self, mode):
""" Set the playback mode.
:param mode: curren tvalid values are "NORMAL" and "SHUFFLE"
"""
self.current_play_mode = getattr(self.play_modes, mode)
self.__update_play_queue_order()
def current_song_title_and_artist(self):
""" Retrieve the current track's title and artist name.
"""
logging.info("current_song_title_and_artist")
song = self.now_playing_song
if song:
title = to_ascii(song.get('title'))
artist = to_ascii(song.get('artist'))
logging.info("Now playing %s by %s", title, artist)
return artist, title
else:
return '', ''
def current_song_album_and_duration(self):
""" Retrieve the current track's album and duration.
"""
logging.info("current_song_album_and_duration")
song = self.now_playing_song
if song:
album = to_ascii(song.get('album'))
duration = to_ascii \
(song.get('durationMillis'))
logging.info("album %s duration %s", album, duration)
return album, int(duration)
else:
return '', 0
def current_track_and_album_total(self):
"""Return the current track number and the total number of tracks in the
album, if known.
"""
logging.info("current_track_and_album_total")
song = self.now_playing_song
track = 0
total = 0
if song:
try:
track = song['trackNumber']
total = song['totalTrackCount']
logging.info("track number %s total tracks %s", track, total)
except KeyError:
logging.info("trackNumber or totalTrackCount : not found")
else:
logging.info("current_song_track_number_"
"and_total_tracks : not found")
return track, total
def current_song_year(self):
""" Return the current track's year of publication.
"""
logging.info("current_song_year")
song = self.now_playing_song
year = 0
if song:
try:
year = song['year']
logging.info("track year %s", year)
except KeyError:
logging.info("year : not found")
else:
logging.info("current_song_year : not found")
return year
def current_song_genre(self):
""" Return the current track's genre.
"""
logging.info("current_song_genre")
song = self.now_playing_song
if song:
genre = to_ascii(song.get('genre'))
logging.info("genre %s", genre)
return genre
else:
return ''
def current_song_album_art(self):
""" Return the current track's album art image.
"""
logging.info("current_song_art")
song = self.now_playing_song
if song:
artref = song.get('albumArtRef')
if artref and len(artref) > 0:
url = to_ascii(artref[0].get('url'))
logging.info("url %s", url)
return url
return ''
def clear_queue(self):
""" Clears the playback queue.
"""
self.queue = list()
self.queue_index = -1
def enqueue_tracks(self, arg):
""" Search the user's library for tracks and add
them to the playback queue.
:param arg: a track search term
"""
try:
songs = self.__gmusic.get_all_songs()
track_hits = list ()
for song in songs:
song_title = song['title']
if arg.lower() in song_title.lower():
track_hits.append(song)
print_nfo("[Google Play Music] [Track] '{0}'." \
.format(to_ascii(song_title)))
if not len(track_hits):
print_wrn("[Google Play Music] '{0}' not found. "\
"Feeling lucky?." \
.format(arg.encode('utf-8')))
random.seed()
track_hits = random.sample(songs, MAX_TRACKS)
for hit in track_hits:
song_title = hit['title']
print_nfo("[Google Play Music] [Track] '{0}'." \
.format(to_ascii(song_title)))
if not len(track_hits):
raise KeyError
tracks_added = self.__enqueue_tracks(track_hits)
logging.info("Added %d tracks from %s to queue", \
tracks_added, arg)
self.__update_play_queue_order()
except KeyError:
raise KeyError("Track not found : {0}".format(arg))
def enqueue_artist(self, arg):
""" Search the user's library for tracks from the given artist and add
them to the playback queue.
:param arg: an artist
"""
try:
self.__update_local_library()
artist = None
artist_dict = None
if arg not in self.library.keys():
for name, art in self.library.iteritems():
if arg.lower() in name.lower():
artist = name
artist_dict = art
if arg.lower() != name.lower():
print_wrn("[Google Play Music] '{0}' not found. " \
"Playing '{1}' instead." \
.format(arg.encode('utf-8'), \
name.encode('utf-8')))
break
if not artist:
# Play some random artist from the library
random.seed()
artist = random.choice(self.library.keys())
artist_dict = self.library[artist]
print_wrn("[Google Play Music] '{0}' not found. "\
"Feeling lucky?." \
.format(arg.encode('utf-8')))
else:
artist = arg
artist_dict = self.library[arg]
tracks_added = 0
for album in artist_dict:
tracks_added += self.__enqueue_tracks(artist_dict[album])
print_wrn("[Google Play Music] Playing '{0}'." \
.format(to_ascii(artist)))
self.__update_play_queue_order()
except KeyError:
raise KeyError("Artist not found : {0}".format(arg))
def enqueue_album(self, arg):
""" Search the user's library for albums with a given name and add
them to the playback queue.
"""
try:
self.__update_local_library()
album = None
artist = None
tentative_album = None
tentative_artist = None
for library_artist in self.library:
for artist_album in self.library[library_artist]:
print_nfo("[Google Play Music] [Album] '{0}'." \
.format(to_ascii(artist_album)))
if not album:
if arg.lower() == artist_album.lower():
album = artist_album
artist = library_artist
break
if not tentative_album:
if arg.lower() in artist_album.lower():
tentative_album = artist_album
tentative_artist = library_artist
if album:
break
if not album and tentative_album:
album = tentative_album
artist = tentative_artist
print_wrn("[Google Play Music] '{0}' not found. " \
"Playing '{1}' instead." \
.format(arg.encode('utf-8'), \
album.encode('utf-8')))
if not album:
# Play some random album from the library
random.seed()
artist = random.choice(self.library.keys())
album = random.choice(self.library[artist].keys())
print_wrn("[Google Play Music] '{0}' not found. "\
"Feeling lucky?." \
.format(arg.encode('utf-8')))
if not album:
raise KeyError("Album not found : {0}".format(arg))
self.__enqueue_tracks(self.library[artist][album])
print_wrn("[Google Play Music] Playing '{0}'." \
.format(to_ascii(album)))
self.__update_play_queue_order()
except KeyError:
raise KeyError("Album not found : {0}".format(arg))
def enqueue_playlist(self, arg):
"""Search the user's library for playlists with a given name
and add the tracks of the first match to the playback queue.
Requires Unlimited subscription.
"""
try:
self.__update_local_library()
self.__update_playlists()
self.__update_playlists_unlimited()
playlist = None
playlist_name = None
for name, plist in self.playlists.items():
print_nfo("[Google Play Music] [Playlist] '{0}'." \
.format(to_ascii(name)))
if arg not in self.playlists.keys():
for name, plist in self.playlists.iteritems():
if arg.lower() in name.lower():
playlist = plist
playlist_name = name
if arg.lower() != name.lower():
print_wrn("[Google Play Music] '{0}' not found. " \
"Playing '{1}' instead." \
.format(arg.encode('utf-8'), \
to_ascii(name)))
break
else:
playlist_name = arg
playlist = self.playlists[arg]
random.seed()
x = 0
while (not playlist or not len(playlist)) and x < 3:
x += 1
# Play some random playlist from the library
playlist_name = random.choice(self.playlists.keys())
playlist = self.playlists[playlist_name]
print_wrn("[Google Play Music] '{0}' not found or found empty. "\
"Feeling lucky?." \
.format(arg.encode('utf-8')))
if not len(playlist):
raise KeyError
self.__enqueue_tracks(playlist)
print_wrn("[Google Play Music] Playing '{0}'." \
.format(to_ascii(playlist_name)))
self.__update_play_queue_order()
except KeyError:
raise KeyError("Playlist not found or found empty : {0}".format(arg))
def enqueue_podcast(self, arg):
"""Search Google Play Music for a podcast series and add its tracks to the
playback queue ().
Requires Unlimited subscription.
"""
print_msg("[Google Play Music] [Retrieving podcasts] : '{0}'. " \
.format(self.__email))
try:
self.__enqueue_podcast(arg)
if not len(self.queue):
raise KeyError
logging.info("Added %d episodes from '%s' to queue", \
len(self.queue), arg)
self.__update_play_queue_order()
except KeyError:
raise KeyError("Podcast not found : {0}".format(arg))
except CallFailure:
raise RuntimeError("Operation requires an Unlimited subscription.")
def enqueue_station_unlimited(self, arg):
"""Search the user's library for a station with a given name
and add its tracks to the playback queue.
Requires Unlimited subscription.
"""
try:
# First try to find a suitable station in the user's library
self.__enqueue_user_station_unlimited(arg)
if not len(self.queue):
# If no suitable station is found in the user's library, then
# search google play unlimited for a potential match.
self.__enqueue_station_unlimited(arg)
if not len(self.queue):
raise KeyError
except KeyError:
raise KeyError("Station not found : {0}".format(arg))
def enqueue_genre_unlimited(self, arg):
"""Search Unlimited for a genre with a given name and add its
tracks to the playback queue.
Requires Unlimited subscription.
"""
print_msg("[Google Play Music] [Retrieving genres] : '{0}'. " \
.format(self.__email))
try:
all_genres = list()
root_genres = self.__gmusic.get_genres()
second_tier_genres = list()
for root_genre in root_genres:
second_tier_genres += self.__gmusic.get_genres(root_genre['id'])
all_genres += root_genres
all_genres += second_tier_genres
for genre in all_genres:
print_nfo("[Google Play Music] [Genre] '{0}'." \
.format(to_ascii(genre['name'])))
genre = dict()
if arg not in all_genres:
genre = next((g for g in all_genres \
if arg.lower() in to_ascii(g['name']).lower()), \
None)
tracks_added = 0
while not tracks_added:
if not genre and len(all_genres):
# Play some random genre from the search results
random.seed()
genre = random.choice(all_genres)
print_wrn("[Google Play Music] '{0}' not found. "\
"Feeling lucky?." \
.format(arg.encode('utf-8')))
genre_name = genre['name']
genre_id = genre['id']
station_id = self.__gmusic.create_station(genre_name, \
None, None, None, genre_id)
num_tracks = MAX_TRACKS
tracks = self.__gmusic.get_station_tracks(station_id, num_tracks)
tracks_added = self.__enqueue_tracks(tracks)
logging.info("Added %d tracks from %s to queue", tracks_added, genre_name)
if not tracks_added:
# This will produce another iteration in the loop
genre = None
print_wrn("[Google Play Music] Playing '{0}'." \
.format(to_ascii(genre['name'])))
self.__update_play_queue_order()
except KeyError:
raise KeyError("Genre not found : {0}".format(arg))
except CallFailure:
raise RuntimeError("Operation requires an Unlimited subscription.")
def enqueue_situation_unlimited(self, arg):
"""Search Unlimited for a situation with a given name and add its
tracks to the playback queue.
Requires Unlimited subscription.
"""
print_msg("[Google Play Music] [Retrieving situations] : '{0}'. " \
.format(self.__email))
try:
self.__enqueue_situation_unlimited(arg)
if not len(self.queue):
raise KeyError
logging.info("Added %d tracks from %s to queue", \
len(self.queue), arg)
except KeyError:
raise KeyError("Situation not found : {0}".format(arg))
except CallFailure:
raise RuntimeError("Operation requires an Unlimited subscription.")
def enqueue_artist_unlimited(self, arg):
"""Search Unlimited for an artist and add the artist's 200 top tracks to the
playback queue.
Requires Unlimited subscription.
"""
try:
artist = self.__gmusic_search(arg, 'artist')
include_albums = False
max_top_tracks = MAX_TRACKS
max_rel_artist = 0
artist_tracks = dict()
if artist:
artist_tracks = self.__gmusic.get_artist_info \
(artist['artist']['artistId'],
include_albums, max_top_tracks,
max_rel_artist)['topTracks']
if not artist_tracks:
raise KeyError
for track in artist_tracks:
song_title = track['title']
print_nfo("[Google Play Music] [Track] '{0}'." \
.format(to_ascii(song_title)))
tracks_added = self.__enqueue_tracks(artist_tracks)
logging.info("Added %d tracks from %s to queue", \
tracks_added, arg)
self.__update_play_queue_order()
except KeyError:
raise KeyError("Artist not found : {0}".format(arg))
except CallFailure:
raise RuntimeError("Operation requires an Unlimited subscription.")
def enqueue_album_unlimited(self, arg):
"""Search Unlimited for an album and add its tracks to the
playback queue.
Requires Unlimited subscription.
"""
try:
album = self.__gmusic_search(arg, 'album')
album_tracks = dict()
if album:
album_tracks = self.__gmusic.get_album_info \
(album['album']['albumId'])['tracks']
if not album_tracks:
raise KeyError
print_wrn("[Google Play Music] Playing '{0}'." \
.format((album['album']['name']).encode('utf-8')))
tracks_added = self.__enqueue_tracks(album_tracks)
logging.info("Added %d tracks from %s to queue", \
tracks_added, arg)
self.__update_play_queue_order()
except KeyError:
raise KeyError("Album not found : {0}".format(arg))
except CallFailure:
raise RuntimeError("Operation requires an Unlimited subscription.")
def enqueue_tracks_unlimited(self, arg):
""" Search Unlimited for a track name and add all the matching tracks
to the playback queue.
Requires Unlimited subscription.
"""
print_msg("[Google Play Music] [Retrieving library] : '{0}'. " \
.format(self.__email))
try:
max_results = MAX_TRACKS
track_hits = self.__gmusic.search(arg, max_results)['song_hits']
if not len(track_hits):
# Do another search with an empty string
track_hits = self.__gmusic.search("", max_results)['song_hits']
print_wrn("[Google Play Music] '{0}' not found. "\
"Feeling lucky?." \
.format(arg.encode('utf-8')))
tracks = list()
for hit in track_hits:
tracks.append(hit['track'])
tracks_added = self.__enqueue_tracks(tracks)
logging.info("Added %d tracks from %s to queue", \
tracks_added, arg)
self.__update_play_queue_order()
except KeyError:
raise KeyError("Playlist not found : {0}".format(arg))
except CallFailure:
raise RuntimeError("Operation requires an Unlimited subscription.")
def enqueue_playlist_unlimited(self, arg):
"""Search Unlimited for a playlist name and add all its tracks to the
playback queue.
Requires Unlimited subscription.
"""
print_msg("[Google Play Music] [Retrieving playlists] : '{0}'. " \
.format(self.__email))
try:
playlist_tracks = list()
playlist_hits = self.__gmusic_search(arg, 'playlist')
if playlist_hits:
playlist = playlist_hits['playlist']
playlist_contents = self.__gmusic.get_shared_playlist_contents(playlist['shareToken'])
else:
raise KeyError
print_nfo("[Google Play Music] [Playlist] '{}'." \
.format(playlist['name']).encode('utf-8'))
for item in playlist_contents:
print_nfo("[Google Play Music] [Playlist Track] '{} by {} (Album: {}, {})'." \
.format((item['track']['title']).encode('utf-8'),
(item['track']['artist']).encode('utf-8'),
(item['track']['album']).encode('utf-8'),
(item['track']['year'])))
track = item['track']
playlist_tracks.append(track)
if not playlist_tracks:
raise KeyError
tracks_added = self.__enqueue_tracks(playlist_tracks)
logging.info("Added %d tracks from %s to queue", \
tracks_added, arg)
self.__update_play_queue_order()
except KeyError:
raise KeyError("Playlist not found : {0}".format(arg))
except CallFailure:
raise RuntimeError("Operation requires an Unlimited subscription.")
def enqueue_promoted_tracks_unlimited(self):
""" Retrieve the url of the next track in the playback queue.
"""
try:
tracks = self.__gmusic.get_promoted_songs()
count = 0
for track in tracks:
store_track = self.__gmusic.get_track_info(track['storeId'])
if u'id' not in store_track.keys():
store_track[u'id'] = store_track['storeId']
self.queue.append(store_track)
count += 1
if count == 0:
print_wrn("[Google Play Music] Operation requires " \
"an Unlimited subscription.")
logging.info("Added %d Unlimited promoted tracks to queue", \
count)
self.__update_play_queue_order()
except CallFailure:
raise RuntimeError("Operation requires an Unlimited subscription.")
def next_url(self):
""" Retrieve the url of the next track in the playback queue.
"""
if len(self.queue):
self.queue_index += 1
if (self.queue_index < len(self.queue)) \
and (self.queue_index >= 0):
next_song = self.queue[self.play_queue_order[self.queue_index]]
return self.__retrieve_track_url(next_song)
else:
self.queue_index = -1
return self.next_url()
else:
return ''
def prev_url(self):
""" Retrieve the url of the previous track in the playback queue.
"""
if len(self.queue):
self.queue_index -= 1
if (self.queue_index < len(self.queue)) \
and (self.queue_index >= 0):
prev_song = self.queue[self.play_queue_order[self.queue_index]]
return self.__retrieve_track_url(prev_song)
else:
self.queue_index = len(self.queue)
return self.prev_url()
else:
return ''
def __update_play_queue_order(self):
""" Update the queue playback order.
A sequential order is applied if the current play mode is "NORMAL" or a
random order if current play mode is "SHUFFLE"
"""
total_tracks = len(self.queue)
if total_tracks:
if not len(self.play_queue_order):
# Create a sequential play order, if empty
self.play_queue_order = range(total_tracks)
if self.current_play_mode == self.play_modes.SHUFFLE:
random.shuffle(self.play_queue_order)
print_nfo("[Google Play Music] [Tracks in queue] '{0}'." \
.format(total_tracks))
def __retrieve_track_url(self, song):
""" Retrieve a song url
"""
if song.get('episodeId'):
song_url = self.__gmusic.get_podcast_episode_stream_url(song['episodeId'], self.__device_id)
else:
song_url = self.__gmusic.get_stream_url(song['id'], self.__device_id)
try:
self.now_playing_song = song
return song_url
except AttributeError:
logging.info("Could not retrieve the song url!")
raise
def __update_local_library(self):
""" Retrieve the songs and albums from the user's library
"""
print_msg("[Google Play Music] [Retrieving library] : '{0}'. " \
.format(self.__email))
songs = self.__gmusic.get_all_songs()
self.playlists[self.thumbs_up_playlist_name] = list()
# Retrieve the user's song library
for song in songs:
if "rating" in song and song['rating'] == "5":
self.playlists[self.thumbs_up_playlist_name].append(song)
song_id = song['id']
song_artist = song['artist']
song_album = song['album']
self.song_map[song_id] = song
if song_artist == "":
song_artist = "Unknown Artist"
if song_album == "":
song_album = "Unknown Album"
if song_artist not in self.library:
self.library[song_artist] = CaseInsensitiveDict()
self.library[song_artist][self.all_songs_album_title] = list()
if song_album not in self.library[song_artist]:
self.library[song_artist][song_album] = list()
self.library[song_artist][song_album].append(song)
self.library[song_artist][self.all_songs_album_title].append(song)
# Sort albums by track number
for artist in self.library.keys():
logging.info("Artist : %s", to_ascii(artist))
for album in self.library[artist].keys():
logging.info(" Album : %s", to_ascii(album))
if album == self.all_songs_album_title:
sorted_album = sorted(self.library[artist][album],
key=lambda k: k['title'])
else:
sorted_album = sorted(self.library[artist][album],
key=lambda k: k.get('trackNumber',
0))
self.library[artist][album] = sorted_album
def __update_stations_unlimited(self):
""" Retrieve stations (Unlimited)
"""
self.stations.clear()
stations = self.__gmusic.get_all_stations()
self.stations[u"I'm Feeling Lucky"] = 'IFL'
for station in stations:
station_name = station['name']
logging.info("station name : %s", to_ascii(station_name))
self.stations[station_name] = station['id']
def __enqueue_user_station_unlimited(self, arg):
""" Enqueue a user station (Unlimited)
"""
print_msg("[Google Play Music] [Station search "\
"in user's library] : '{0}'. " \
.format(self.__email))
self.__update_stations_unlimited()
station_name = arg
station_id = None
for name, st_id in self.stations.iteritems():
print_nfo("[Google Play Music] [Station] '{0}'." \
.format(to_ascii(name)))
if arg not in self.stations.keys():
for name, st_id in self.stations.iteritems():
if arg.lower() in name.lower():
station_id = st_id
station_name = name
break
else:
station_id = self.stations[arg]
num_tracks = MAX_TRACKS
tracks = list()
if station_id:
try:
tracks = self.__gmusic.get_station_tracks(station_id, \
num_tracks)
except KeyError:
raise RuntimeError("Operation requires an "
"Unlimited subscription.")
tracks_added = self.__enqueue_tracks(tracks)
if tracks_added:
if arg.lower() != station_name.lower():
print_wrn("[Google Play Music] '{0}' not found. " \
"Playing '{1}' instead." \
.format(arg.encode('utf-8'), name.encode('utf-8')))
logging.info("Added %d tracks from %s to queue", tracks_added, arg)
self.__update_play_queue_order()
else:
print_wrn("[Google Play Music] '{0}' has no tracks. " \
.format(station_name))
if not len(self.queue):
print_wrn("[Google Play Music] '{0}' " \
"not found in the user's library. " \
.format(arg.encode('utf-8')))
def __enqueue_station_unlimited(self, arg, max_results=MAX_TRACKS, quiet=False):
"""Search for a station and enqueue all of its tracks (Unlimited)
"""
if not quiet:
print_msg("[Google Play Music] [Station search in "\
"Google Play Music] : '{0}'. " \
.format(arg.encode('utf-8')))
try:
station_name = arg
station_id = None
station = self.__gmusic_search(arg, 'station', max_results, quiet)
if station:
station = station['station']
station_name = station['name']
seed = station['seed']
seed_type = seed['seedType']
track_id = seed['trackId'] if seed_type == u'2' else None
artist_id = seed['artistId'] if seed_type == u'3' else None
album_id = seed['albumId'] if seed_type == u'4' else None
genre_id = seed['genreId'] if seed_type == u'5' else None
playlist_token = seed['playlistShareToken'] if seed_type == u'8' else None
curated_station_id = seed['curatedStationId'] if seed_type == u'9' else None
num_tracks = max_results
tracks = list()
try:
station_id \
= self.__gmusic.create_station(station_name, \
track_id, \
artist_id, \
album_id, \
genre_id, \
playlist_token, \
curated_station_id)
tracks \
= self.__gmusic.get_station_tracks(station_id, \
num_tracks)
except KeyError:
raise RuntimeError("Operation requires an "
"Unlimited subscription.")
tracks_added = self.__enqueue_tracks(tracks)
if tracks_added:
if not quiet:
print_wrn("[Google Play Music] [Station] : '{0}'." \
.format(station_name.encode('utf-8')))
logging.info("Added %d tracks from %s to queue", \
tracks_added, arg.encode('utf-8'))
self.__update_play_queue_order()
except KeyError:
raise KeyError("Station not found : {0}".format(arg))
def __enqueue_situation_unlimited(self, arg):
"""Search for a situation and enqueue all of its tracks (Unlimited)
"""
print_msg("[Google Play Music] [Situation search in "\
"Google Play Music] : '{0}'. " \
.format(arg.encode('utf-8')))
try:
situation_hits = self.__gmusic.search(arg)['situation_hits']
# If the search didn't return results, just do another search with
# an empty string
if not len(situation_hits):
situation_hits = self.__gmusic.search("")['situation_hits']
print_wrn("[Google Play Music] '{0}' not found. "\
"Feeling lucky?." \
.format(arg.encode('utf-8')))
# Try to find a "best result", if one exists
situation = next((hit for hit in situation_hits \
if 'best_result' in hit.keys() \
and hit['best_result'] == True), None)
num_tracks = MAX_TRACKS
# If there is no best result, then get a selection of tracks from
# each situation. At least we'll play some music.
if not situation and len(situation_hits):
max_results = num_tracks / len(situation_hits)
for hit in situation_hits:
situation = hit['situation']
print_nfo("[Google Play Music] [Situation] '{0} : {1}'." \
.format((hit['situation']['title']).encode('utf-8'),
(hit['situation']['description']).encode('utf-8')))
self.__enqueue_station_unlimited(situation['title'], max_results, True)
elif situation:
# There is at list one sitution, enqueue its tracks.
situation = situation['situation']
max_results = num_tracks
self.__enqueue_station_unlimited(situation['title'], max_results, True)
if not situation:
raise KeyError
except KeyError:
raise KeyError("Situation not found : {0}".format(arg))
def __enqueue_podcast(self, arg):
"""Search for a podcast series and enqueue all of its tracks.
"""
print_msg("[Google Play Music] [Podcast search in "\
"Google Play Music] : '{0}'. " \
.format(arg.encode('utf-8')))
try:
podcast_hits = self.__gmusic_search(arg, 'podcast', 10, quiet=False)
if not podcast_hits:
print_wrn("[Google Play Music] [Podcast] 'Search returned zero results'.")
print_wrn("[Google Play Music] [Podcast] 'Are you in a supported region "
"(currently only US and Canada) ?'")
# Use the first podcast retrieved. At least we'll play something.
podcast = dict ()
if podcast_hits and len(podcast_hits):
podcast = podcast_hits['series']
episodes_added = 0
if podcast:
# There is a podcast, enqueue its episodes.
print_nfo("[Google Play Music] [Podcast] 'Playing '{0}' by {1}'." \
.format((podcast['title']).encode('utf-8'),
(podcast['author']).encode('utf-8')))
print_nfo("[Google Play Music] [Podcast] '{0}'." \
.format((podcast['description'][0:150]).encode('utf-8')))
series = self.__gmusic.get_podcast_series_info(podcast['seriesId'])
episodes = series['episodes']
for episode in episodes:
print_nfo("[Google Play Music] [Podcast Episode] '{0} : {1}'." \
.format((episode['title']).encode('utf-8'),
(episode['description'][0:80]).encode('utf-8')))
episodes_added = self.__enqueue_tracks(episodes)
if not podcast or not episodes_added:
raise KeyError
except KeyError:
raise KeyError("Podcast not found or no episodes found: {0}".format(arg))
def __enqueue_tracks(self, tracks):
""" Add tracks to the playback queue
"""
count = 0
for track in tracks:
if u'id' not in track.keys() and track.get('storeId'):
track[u'id'] = track['storeId']
self.queue.append(track)
count += 1
return count
def __update_playlists(self):
""" Retrieve the user's playlists
"""
plists = self.__gmusic.get_all_user_playlist_contents()
for plist in plists:
plist_name = plist.get('name')
tracks = plist.get('tracks')
if plist_name and tracks:
logging.info("playlist name : %s", to_ascii(plist_name))
tracks.sort(key=itemgetter('creationTimestamp'))
self.playlists[plist_name] = list()
for track in tracks:
song_id = track.get('trackId')
if song_id:
song = self.song_map.get(song_id)
if song:
self.playlists[plist_name].append(song)
def __update_playlists_unlimited(self):
""" Retrieve shared playlists (Unlimited)
"""
plists_subscribed_to = [p for p in self.__gmusic.get_all_playlists() \
if p.get('type') == 'SHARED']
for plist in plists_subscribed_to:
share_tok = plist['shareToken']
playlist_items \
= self.__gmusic.get_shared_playlist_contents(share_tok)
plist_name = plist['name']
logging.info("shared playlist name : %s", to_ascii(plist_name))
self.playlists[plist_name] = list()
for item in playlist_items:
try:
song = item['track']
song['id'] = item['trackId']
self.playlists[plist_name].append(song)
except IndexError:
pass
def __gmusic_search(self, query, query_type, max_results=MAX_TRACKS, quiet=False):
""" Search Google Play (Unlimited)
"""
search_results = self.__gmusic.search(query, max_results)[query_type + '_hits']
# This is a workaround. Some podcast results come without these two
# keys in the dictionary
if query_type == "podcast" and len(search_results) \
and not search_results[0].get('navigational_result'):
for res in search_results:
res[u'best_result'] = False
res[u'navigational_result'] = False
res[query_type] = res['series']
result = ''
if query_type != "playlist":
result = next((hit for hit in search_results \
if 'best_result' in hit.keys() \
and hit['best_result'] == True), None)
if not result and len(search_results):
secondary_hit = None
for hit in search_results:
name = ''
if hit[query_type].get('name'):
name = hit[query_type].get('name')
elif hit[query_type].get('title'):
name = hit[query_type].get('title')
if not quiet:
print_nfo("[Google Play Music] [{0}] '{1}'." \
.format(query_type.capitalize(),
(name).encode('utf-8')))
if query.lower() == \
to_ascii(name).lower():
result = hit
break
if query.lower() in \
to_ascii(name).lower():
secondary_hit = hit
if not result and secondary_hit:
result = secondary_hit
if not result and not len(search_results):
# Do another search with an empty string
search_results = self.__gmusic.search("")[query_type + '_hits']
if not result and len(search_results):
# Play some random result from the search results
random.seed()
result = random.choice(search_results)
if not quiet:
print_wrn("[Google Play Music] '{0}' not found. "\
"Feeling lucky?." \
.format(query.encode('utf-8')))
return result