Permalink
Fetching contributors…
Cannot retrieve contributors at this time
157 lines (138 sloc) 6.02 KB
# 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)