Permalink
Cannot retrieve contributors at this time
# Copyright 2016, Cossack Labs Limited | |
# | |
# Licensed under the Apache License, Version 2.0 (the "License"); | |
# you may not use this file except in compliance with the License. | |
# You may obtain a copy of the License at | |
# | |
# http://www.apache.org/licenses/LICENSE-2.0 | |
# | |
# Unless required by applicable law or agreed to in writing, software | |
# distributed under the License is distributed on an "AS IS" BASIS, | |
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
# See the License for the specific language governing permissions and | |
# limitations under the License. | |
# coding: utf-8 | |
import os | |
import argparse | |
import json | |
from base64 import b64decode | |
from random import randint, choice | |
try: | |
# For Python 3.0 and later | |
from urllib.request import urlopen | |
except ImportError: | |
# Fall back to Python 2's urllib2 | |
from urllib2 import urlopen | |
from sqlalchemy import (Table, Column, Integer, MetaData, create_engine, | |
select, Binary, Text, cast) | |
from sqlalchemy.dialects.postgresql import BYTEA | |
from acrawriter import create_acrastruct | |
def get_zone(): | |
"""make http response to AcraServer api to generate new zone and return tuple | |
of zone id and public key | |
""" | |
response = urlopen('{}/getNewZone'.format(ACRA_CONNECTOR_API_ADDRESS)) | |
json_data = response.read().decode('utf-8') | |
zone_data = json.loads(json_data) | |
return zone_data['id'], b64decode(zone_data['public_key']) | |
def get_default(name, value): | |
"""return value from environment variables with name EXAMPLE_<name> | |
or value""" | |
return os.environ.get('EXAMPLE_{}'.format(name.upper()), value) | |
def print_data(zone_id, connection): | |
"""fetch data from database (use zone_id if not empty/None) and print to | |
console""" | |
result = connection.execute( | |
# explicitly pass zone id before related data | |
select([cast(zone_id.encode('utf-8'), BYTEA), test_table])) | |
result = result.fetchall() | |
ZONE_ID_INDEX = 0 | |
print("use zone_id: ", zone_id) | |
print("{:<3} - {} - {} - {}".format("id", 'zone', "data", "raw_data")) | |
for row in result: | |
print( | |
"{:<3} - {} - {} - {}\n".format( | |
row['id'], row[ZONE_ID_INDEX].decode('utf-8'), | |
row['data'].decode('utf-8', errors='ignore'), row['raw_data'])) | |
def write_data(data, connection): | |
zone_id, key = get_zone() | |
print("data: {}\nzone: {}".format(data, zone_id)) | |
# here we encrypt our data and wrap into AcraStruct | |
encrypted_data = create_acrastruct( | |
data.encode('utf-8'), key, zone_id.encode('utf-8')) | |
connection.execute( | |
test_table.insert(), data=encrypted_data, | |
zone_id=zone_id.encode('utf-8'), | |
raw_data=data) | |
if __name__ == '__main__': | |
parser = argparse.ArgumentParser() | |
parser.add_argument('--zone_id', type=str, | |
default=get_default('zone_id', ''), | |
help='Zone id for fetching data. Use only with --print ' | |
'option') | |
parser.add_argument('--db_name', type=str, | |
default=get_default('db_name', 'acra'), | |
help='Database name') | |
parser.add_argument('--db_user', type=str, | |
default=get_default('db_user','test'), | |
help='Database user') | |
parser.add_argument('--db_password', type=str, | |
default=get_default('db_password', 'test'), | |
help='Database user\'s password') | |
parser.add_argument('--port', type=int, | |
default=get_default('port', 9494), | |
help='Port of database or AcraConnector') | |
parser.add_argument('--host', type=str, | |
default=get_default('host', 'localhost'), | |
help='Host of database or AcraConnector') | |
parser.add_argument('--data', type=str, | |
default=get_default('data', ''), | |
help='data to save in ascii. default random data') | |
parser.add_argument('--print', action='store_true', | |
default=get_default('print', False), | |
help='Print data (use --zone_id to set specific ZoneId ' | |
'which will be used to fetch data)') | |
parser.add_argument('-v', '--verbose', dest='verbose', action='store_true', | |
default=get_default('verbose', False), help='verbose') | |
parser.add_argument('--postgresql', action='store_true', | |
default=get_default('postgresql', False), | |
help="Use postgresql driver (default if nothing else " | |
"set)") | |
parser.add_argument('--mysql', action='store_true', | |
default=get_default('mysql', False), | |
help="Use mysql driver") | |
args = parser.parse_args() | |
ACRA_CONNECTOR_API_ADDRESS = get_default( | |
'acra_connector_api_address', 'http://127.0.0.1:9191') | |
# default driver | |
driver = 'postgresql' | |
if args.mysql: | |
driver = 'mysql+pymysql' | |
metadata = MetaData() | |
test_table = Table( | |
'test_example_with_zone', metadata, | |
Column('id', Integer, primary_key=True, nullable=False), | |
Column('zone_id', Binary, nullable=True), | |
Column('data', Binary, nullable=False), | |
Column('raw_data', Text, nullable=False), | |
) | |
engine = create_engine( | |
'{}://{}:{}@{}:{}/{}'.format( | |
driver, args.db_user, args.db_password, args.host, args.port, | |
args.db_name), | |
echo=bool(args.verbose)) | |
connection = engine.connect() | |
metadata.create_all(engine) | |
print('DB driver: {}'.format(driver)) | |
if args.print: | |
print_data(args.zone_id, connection) | |
elif args.data: | |
if args.zone_id: | |
print("To encrypt data script will generate new zone and print " | |
"zone id with public key after execution. Don't use " | |
"--zone_id option with --data option.") | |
exit(1) | |
write_data(args.data, connection) | |
else: | |
print('Use --print or --data options') | |
exit(1) |