mirror of
https://github.com/klzgrad/naiveproxy.git
synced 2024-11-28 16:26:10 +03:00
90 lines
2.8 KiB
Python
90 lines
2.8 KiB
Python
# Copyright 2017 The Chromium Authors. All rights reserved.
|
|
# Use of this source code is governed by a BSD-style license that can be
|
|
# found in the LICENSE file.
|
|
|
|
import hashlib
|
|
import os
|
|
|
|
from pylib.base import output_manager
|
|
from pylib.output import noop_output_manager
|
|
from pylib.utils import logdog_helper
|
|
from pylib.utils import google_storage_helper
|
|
|
|
|
|
class RemoteOutputManager(output_manager.OutputManager):
|
|
|
|
def __init__(self, bucket):
|
|
"""Uploads output files to Google Storage or LogDog.
|
|
|
|
Files will either be uploaded directly to Google Storage or LogDog
|
|
depending on the datatype.
|
|
|
|
Args
|
|
bucket: Bucket to use when saving to Google Storage.
|
|
"""
|
|
super(RemoteOutputManager, self).__init__()
|
|
self._bucket = bucket
|
|
|
|
#override
|
|
def _CreateArchivedFile(self, out_filename, out_subdir, datatype):
|
|
if datatype == output_manager.Datatype.TEXT:
|
|
try:
|
|
logdog_helper.get_logdog_client()
|
|
return LogdogArchivedFile(out_filename, out_subdir, datatype)
|
|
except RuntimeError:
|
|
return noop_output_manager.NoopArchivedFile()
|
|
else:
|
|
if self._bucket is None:
|
|
return noop_output_manager.NoopArchivedFile()
|
|
return GoogleStorageArchivedFile(
|
|
out_filename, out_subdir, datatype, self._bucket)
|
|
|
|
|
|
class LogdogArchivedFile(output_manager.ArchivedFile):
|
|
|
|
def __init__(self, out_filename, out_subdir, datatype):
|
|
super(LogdogArchivedFile, self).__init__(out_filename, out_subdir, datatype)
|
|
self._stream_name = '%s_%s' % (out_subdir, out_filename)
|
|
|
|
def _Link(self):
|
|
return logdog_helper.get_viewer_url(self._stream_name)
|
|
|
|
def _Archive(self):
|
|
with open(self.name, 'r') as f:
|
|
logdog_helper.text(self._stream_name, f.read())
|
|
|
|
|
|
class GoogleStorageArchivedFile(output_manager.ArchivedFile):
|
|
|
|
def __init__(self, out_filename, out_subdir, datatype, bucket):
|
|
super(GoogleStorageArchivedFile, self).__init__(
|
|
out_filename, out_subdir, datatype)
|
|
self._bucket = bucket
|
|
self._upload_path = None
|
|
self._content_addressed = None
|
|
|
|
def _PrepareArchive(self):
|
|
self._content_addressed = (self._datatype in (
|
|
output_manager.Datatype.HTML,
|
|
output_manager.Datatype.PNG,
|
|
output_manager.Datatype.JSON))
|
|
if self._content_addressed:
|
|
sha1 = hashlib.sha1()
|
|
with open(self.name, 'rb') as f:
|
|
sha1.update(f.read())
|
|
self._upload_path = sha1.hexdigest()
|
|
else:
|
|
self._upload_path = os.path.join(self._out_subdir, self._out_filename)
|
|
|
|
def _Link(self):
|
|
return google_storage_helper.get_url_link(
|
|
self._upload_path, self._bucket)
|
|
|
|
def _Archive(self):
|
|
if (self._content_addressed and
|
|
google_storage_helper.exists(self._upload_path, self._bucket)):
|
|
return
|
|
|
|
google_storage_helper.upload(
|
|
self._upload_path, self.name, self._bucket, content_type=self._datatype)
|