How to use adapter_probe method in dbt-osmosis

Best Python code snippet using dbt-osmosis_python

osmosis.py

Source:osmosis.py

...285        except Exception as query_exc:286            raise RuntimeException("Could not connect to Database") from query_exc287        else:288            return adapter289    def adapter_probe(self) -> bool:290        """Check adapter connection, useful for long running processes such as the server or workbench"""291        if not hasattr(self, "adapter") or self.adapter is None:292            return False293        try:294            with self.adapter.connection_named("osmosis-heartbeat"):295                self.adapter.debug_query()296        except Exception:297            # TODO: we can decide to reinit the Adapter here298            return False299        logger().info("Heartbeat received for %s", self.project_name)300        return True301    def fn_threaded_conn(self, fn: Callable[..., T], *args, **kwargs) -> Callable[..., T]:302        """Used for jobs which are intended to be submitted to a thread pool,303        the 'master' thread should always have an available connection for the duration of...

production_line_tool.py

Source:production_line_tool.py

1#!/usr/bin/env python323# Copyright 2022 Silicon Laboratories Inc. www.silabs.com4#5# SPDX-License-Identifier: Zlib6#7# The licensor of this software is Silicon Laboratories Inc.8#9# This software is provided 'as-is', without any express or implied10# warranty. In no event will the authors be held liable for any damages11# arising from the use of this software.12#13# Permission is granted to anyone to use this software for any purpose,14# including commercial applications, and to alter it and redistribute it15# freely, subject to the following restrictions:16#17# 1. The origin of this software must not be misrepresented; you must not18#    claim that you wrote the original software. If you use this software19#    in a product, an acknowledgment in the product documentation would be20#    appreciated but is not required.21# 2. Altered source versions must be plainly marked as such, and must not be22#    misrepresented as being the original software.23# 3. This notice may not be removed or altered from any source distribution.2425'''Certificate Authority2627Handles signing Certificate Signing Requests and retrieving the Static 28Authentication Data from connected embedded devices.2930Prerequisites:31A connected device with the prepared Certificate Signing Request.32Simplicity Commander added to path.33Cryptography module installed.34Certificate Authority created using the create_authority_certificate.py script.35'''36# Metadata37__author__ = 'Silicon Laboratories, Inc'38__copyright__ = 'Copyright 2022, Silicon Laboratories, Inc.'3940import os, sys, pathlib41import datetime42import re43import binascii, struct44import argparse45import subprocess46from cryptography import x50947from cryptography.hazmat import primitives48from cryptography.hazmat.primitives import serialization4950RAM_DATA_LEN_MAX        = 102451STATIC_AUTH_DATA_LEN    = 325253NVM3_CONTROL_BLOCK_KEY = { 'ble': 0x40400, 'btmesh': 0x60400 }54NVM3_CONTROL_BLOCK_SIZE = 1755CHAIN_LINK_DATA_LEN     = 1925657# Control Block bitmap58CERTIFICATE_ON_DEVICE_BIT       = 059DEVICE_EC_KEY_BIT               = 160STATIC_AUTH_DATA_BIT            = 26162def main(level, validity, serial, ip, protocol):63    # Check the presence of Simplicity Commander.64    try:65        subprocess.run(['commander', '-v'], stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT) # Supress output66    except FileNotFoundError:67        raise FileNotFoundError('Cannot find Simplicity Commander.')6869    # Check arguments.70    if level < 0:71        raise Exception('Authority level should be greater than or equal to zero!')72    if validity < 1:73        raise Exception('Validity should be greater than or equal to 1 [days]!')74    if serial is not None and ip is not None:75        raise Exception('At most one of [J-Link serial number] or [IP address] shall be defined.')76    if protocol not in NVM3_CONTROL_BLOCK_KEY:77        raise Exception('Invalid protocol: ' + protocol + ' (choose from: ' + str(NVM3_CONTROL_BLOCK_KEY.keys()) + ')')78    79    # Check connected device.80    cmd = ('commander adapter probe ' + device_argument()).split()81    adapter_probe = str(subprocess_call(cmd).stdout)8283    # Check chip family.84    substr_start = adapter_probe.lower().find('efr')85    substr_end = adapter_probe.find(' ', substr_start)86    family = adapter_probe[substr_start:substr_end]87    print(family, 'detected.\n\n')8889    if 'xg22' in family.lower():90        ram_start = 0x20000000 # The CSR allocation is different for xG22 devices.91    else:92        ram_start = 0x2001000093    94    # Paths95    if level == 0:96        path_auth_dir = os.path.join(os.path.abspath(os.path.join((__file__), '..')), 'central_authority')97    else:98        path_auth_dir = os.path.join(os.path.abspath(os.path.join((__file__), '..')), 'intermediate_' + str(level) + '_authority')99    100    if not os.path.exists(path_auth_dir):101        raise FileNotFoundError(path_auth_dir + ' authority is incomplete.')102    103    # Retrieve RAM data.104    ram_data = read_ram(ram_start, RAM_DATA_LEN_MAX)105106    # Retrieve Control Block from NVM3.107    path_nvm3_dump = dump_nvm3(path_auth_dir)108    control_block  = ControlBlock(path_nvm3_dump)109110    if control_block.static_auth_present == True:111        # Get Static Authentication Data from RAM data.112        static_auth = get_static_auth(ram_data) # static_auth is ready for further processing.113114    if control_block.device_ec_key_present == False:115        print('EC key pair is missing from the device. It is needed to generate a CSR. Exiting.')116        os.remove(path_nvm3_dump)117        sys.exit(0)118119    # Get Certificate Signing Request from RAM data.120    path_csr = get_csr(ram_data, path_auth_dir)121122    # Sign it and create the Device Certificate.123    path_cert_pem = sign_csr(path_auth_dir, path_csr, validity)124    os.remove(path_csr) # Remove CSR file since it is not needed anymore.125126    # Exit if the device is not configured to hold the certificate.127    if control_block.certificate_on_device == False:128        print('The device is not configured to hold the certificate.')129        print(path_cert_pem, ' is ready for further processing.')130        os.remove(path_nvm3_dump)131        sys.exit(0)132133    # Change the format of a certificate from PEM to DER.134    path_cert_der = convert_crt(path_cert_pem)135136    # Add the signed certificate to the NVM3 image according to the Control Block.137    patch_nvm3(path_nvm3_dump,138               path_cert_der,139               control_block.nvm3_key[CERTIFICATE_ON_DEVICE_BIT],140               control_block.max_link_data_len)141142    # Flash back the modified NVM3 content onto the device.143    flash_nvm3(path_nvm3_dump)144145    # Clean up.146    os.remove(path_cert_der)147    os.remove(path_nvm3_dump)148149def read_ram(start_address, memory_size):150    '''Get data from the RAM in the given range.151152    Keyword arguments:153    start_address -- Address to start reading from.154    memory_size -- Memory size to read.155    Return values:156    The retrieved data from the RAM.157    '''158    cmd = ('commander readmem '159           + device_argument()160           + ' --range ' + hex(start_address) + ':+' + hex(memory_size)).split()161    output = subprocess_call(cmd).stdout162163    # Grab raw bytes from the output.164    output = re.findall(r'^([0-9A-Fa-f]{8}):(\s--)*((\s[0-9A-Fa-f]{2}){1,16})(\s--)*',165                        output.decode('utf-8'),166                        re.MULTILINE)167    s = ''168169    for item in output:170      s = s + item[2].replace(' ', '')171172    return binascii.unhexlify(s)173174def get_static_auth(ram_data):175    '''Static Authentication Data.176177    Keyword arguments:178    ram_data -- Data retrieved from the RAM using read_ram179    Return values:180    static_auth -- Static Authentication Data.181    '''182    # Select static authentication data.183    static_auth = ram_data[1:STATIC_AUTH_DATA_LEN]184185    print('Static Authentication Data has been successfully retrieved from the device.\n\n')186    return static_auth187188def get_csr(ram_data, path_dir):189    '''Get Certificate Signing Request.190191    Keyword arguments:192    ram_data -- Data retrieved from the RAM using read_ram193    path_dir -- Path to working directory where the CSR file shall be created.194    Return values:195    path_csr -- Path to the Certificate Signing Request file.196    '''197    # Select the certificate signing request.198    is_complete = ram_data[0] == 1 # Completion bit199    tuple = struct.unpack('<H', ram_data[STATIC_AUTH_DATA_LEN + 1 : STATIC_AUTH_DATA_LEN + 3])200    csr_len = tuple[0]201202    if is_complete == False or csr_len == 0:203        raise Exception('Certificate signing request generation is incomplete!')204205    csr_bytes = ram_data[STATIC_AUTH_DATA_LEN + 3 : STATIC_AUTH_DATA_LEN + 3 + csr_len]206    csr = x509.load_der_x509_csr(csr_bytes)207208    # Write CSR to file.209    common_name = csr.subject.get_attributes_for_oid(x509.oid.NameOID.COMMON_NAME)[0].value210    path_csr = os.path.join(path_dir, common_name + '.csr')211    212    with open(path_csr, 'wb') as f:213        f.write(csr.public_bytes(serialization.Encoding.PEM))214215    print('Certificate Signing Request has been successfully retrieved from the device.\n\n')216    return path_csr217218def sign_csr(path_auth_dir, path_device_csr, validity):219    '''Create a device certificate by signing a Certificate Signing Request.220    221    Keyword arguments:222    path_auth_dir -- Path to the certificate authority that shall sign the request.223    path_device_csr -- The location of the certificate signing request file in 224                       PEM format.225    validity -- The valid period of the certificate in days starting226                from the moment of signing. (Default: 365)227    Return values:228    path_device_cert -- The location of the certificate file in PEM format.229    '''230231    # Set paths and check prerequisites.232    path_auth_key = os.path.join(path_auth_dir, 'private_key.pem')233    path_auth_cert = os.path.join(path_auth_dir, 'certificate.pem')234235    if not os.path.exists(path_auth_key) or not os.path.exists(path_auth_cert):236        raise FileNotFoundError(path_auth_dir + ' authority is incomplete.')237238    path_device_csr = pathlib.Path(path_device_csr)239240    if not os.path.exists(path_device_csr):241        raise FileNotFoundError('Certificate signing request cannot be found at ' + path_device_csr)242243    path_device_cert = pathlib.Path(os.path.splitext(path_device_csr)[0] + '.crt')244245    # Load certificate signing request246    with open(path_device_csr, 'rb') as f:247        csr = x509.load_pem_x509_csr(f.read())248249    # Load private key and certificate of the authority.250    with open(path_auth_key, 'rb') as f:251        authority_key = primitives.serialization.load_pem_private_key(f.read(), password=None)252    253    with open(path_auth_cert, 'rb') as f:254        authority_cert = x509.load_pem_x509_certificate(f.read())255256    # Check the validity of the higher authority.257    # Note: if there is a revocation list, it should be also checked if the certificate is revoked or not.258    if datetime.datetime.utcnow() < authority_cert.not_valid_before or authority_cert.not_valid_after < datetime.datetime.utcnow():259        raise Exception('The validity period of ' + path_auth_cert + ' has expired.')260    261    # Check if the CSR signature is valid262    if not csr.is_signature_valid:263        raise Exception('Invalid CSR signature!')264    265    # Create a new certificate from the CSR.266    cert = (267        x509.CertificateBuilder()268            .subject_name(csr.subject)269            .issuer_name(authority_cert.issuer)270            .public_key(csr.public_key())271            .serial_number(x509.random_serial_number())272            .not_valid_before(datetime.datetime.utcnow())273            .not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(validity))274    )275276    # Add requested extensions277    for ext in csr.extensions:278        cert = cert.add_extension(ext.value, ext.critical)279    280    # Sign certificate with the private key of the authority.281    cert = cert.sign(authority_key, csr.signature_hash_algorithm)282283    # Verify signature.284    authority_cert.public_key().verify(cert.signature,285                                       cert.tbs_certificate_bytes,286                                       primitives.asymmetric.ec.ECDSA(csr.signature_hash_algorithm))287    288    # Write the signed certificate to file in PEM format.289    with open(path_device_cert, 'wb') as f:290        f.write(cert.public_bytes(primitives.serialization.Encoding.PEM))291    292    print('Signing completed.\n' + str(path_device_cert) + ' created.\n\n')293    return path_device_cert294295def convert_crt(path_pem):296    '''Make a copy of a PEM certificate file in DER format.297298    Keyword arguments:299    path_pem -- Path to the certificate file in PEM format.300    Return values:301    path_der -- Path to the formatted certificate file.302    '''303    # Check input304    if not os.path.exists(path_pem):305        raise FileNotFoundError('Cannot find certificate file.')306307    # Load certificate308    with open(path_pem, 'rb') as f:309        crt = x509.load_pem_x509_certificate(f.read())310    311    # Write the certificate to file in DER format.312    path_der = pathlib.Path(os.path.splitext(path_pem)[0] + '.der')313314    with open(path_der, 'wb') as f:315        f.write(crt.public_bytes(primitives.serialization.Encoding.DER))316    317    print('Certificate converted from PEM to DER successfully. ' + str(path_der) + ' created.')318    return path_der319320def dump_nvm3(path_dir):321    '''Retrieve NVM3 content in s37 format.322323    Keyword arguments:324    path_dir -- Path to working directory where the NVM3 image file shall be created.325    Return values:326    path_nvm3_dump -- Path to the NVM3 image file.327    '''328    path_nvm3_dump = os.path.join(path_dir, 'nvm3_dump.s37')329330    # Remove temp file if exist331    if os.path.exists(path_nvm3_dump):332        print(path_nvm3_dump + ' already exists. Removing.\n')333        os.remove(path_nvm3_dump)334    335    # Search for the NVM3 area in the device's flash and dump the content to file.336    subprocess_call(('commander nvm3 read ' + device_argument() + ' -o ' + path_nvm3_dump).split())337    print('NVM3 dumped to ' + str(path_nvm3_dump))338    return path_nvm3_dump339340def patch_nvm3(path_nvm3_dump, path_data, nvm3_start_key, nvm3_object_size):341    '''Create an NVM3 patch file containing the given data and apply on the NVM3 image.342343    Keyword arguments:344    path_nvm3_dump -- Path to the NVM3 image.345    path_data -- A binary file. The data of this file will be added to the NVM3 image.346    nvm3_start_key -- NVM3 objects will be created or overwritten starting from this key.347    nvm3_object_size -- The size of each NVM3 object data.348    '''349350    path_nvm3_patch = str(pathlib.Path(os.path.splitext(path_nvm3_dump)[0] + '.patch'))351    nvm3_key = nvm3_start_key352353    with open(path_data, 'rb') as f:354        crt = f.read()355356    # Create patch file357    with open(path_nvm3_patch, 'w') as f:358        while nvm3_key != 0:359            # Select next chunk.360            if len(crt) >= nvm3_object_size:361                length = nvm3_object_size362            else:363                length = len(crt)364365            chunk = crt[:length]366            crt = crt[length:]367368            if len(crt) > 0:369                next_nvm3_key = nvm3_key + 1370                header = 0x0001 | 2 | next_nvm3_key << 2371            else: # The last chunk is being processed.372                next_nvm3_key = 0373                header = 0x0001374            375            # Construct next line and write to file.376            key = NVM3_CONTROL_BLOCK_KEY[args.protocol] | nvm3_key # Add key prefix.377            data = struct.pack(f'<HHH', header, 0x0001, length) + chunk # Pack data chunk.378            f.write(hex(key) + ' : OBJ : ' + binascii.hexlify(data).decode('utf-8') + '\n') # <key>:<type>:<data>379380            nvm3_key = next_nvm3_key381382    # Apply patch383    subprocess_call(['commander', 'nvm3', 'set', path_nvm3_dump, '--nvm3file', path_nvm3_patch, '--outfile', path_nvm3_dump])384    print(str(path_nvm3_dump) + ' patched successfully.')385    os.remove(path_nvm3_patch)386387def flash_nvm3(path_nvm3_dump):388    '''Flash NVM3 content.389390    Keyword arguments:391    path_nvm3_dump -- Path to the NVM3 image.392    '''393    subprocess_call(('commander flash ' + path_nvm3_dump + ' ' + device_argument()).split())394    print('The device is updated with the new Flash image.')395396class ControlBlock():397    def __init__(self, path_nvm3_dump):398        '''Get Control Block from an NVM3 image file.399400        Keyword arguments:401        path_nvm3_dump -- Path to the NVM3 image.402        Return values:403        ControlBlock -- The retrieved and parsed Control Block.404        '''405        nvm3_object = self.get_nvm3_object(path_nvm3_dump)406        self.parse_data(nvm3_object)407        self.validate()408    409    def get_nvm3_object(self, path_nvm3_dump):410        # Read NVM3 object from the NVM3 image that belongs to the control block key.411        cmd = ('commander nvm3 parse ' + path_nvm3_dump + ' --key ' + str(NVM3_CONTROL_BLOCK_KEY[args.protocol])).split()412        nvm3_object = subprocess_call(cmd).stdout.decode('utf-8')413414        # Check Control Block key.415        key = re.findall(r'^Key\s+:\s+(0x\d+).*$', nvm3_object, re.MULTILINE)416417        if len(re.findall(r'^Found NVM3 range:', nvm3_object, re.MULTILINE)) == 0 or len(key) == 0:418            raise Exception('NVM3 range cannot be found.')419420        self.key = key[0]421        print('Control Block found at ' + self.key)422        return nvm3_object423    424    def parse_data(self, nvm3_object):425        # Get raw bytes of the Control Block.426        data = ''427428        for item in re.findall(r'^([0-9A-Fa-f]{8}):((\s[0-9A-Fa-f]{2}){1,16})', nvm3_object, re.MULTILINE):429            data = data + item[1].replace(' ', '')430431        data = binascii.unhexlify(data)432433        if data == None or len(data) != NVM3_CONTROL_BLOCK_SIZE:434            raise Exception('Invalid Control Block!')435436        # Unpack the data of the NVM3 object.437        self.nvm3_key = [0, 0, 0]438        (self.header,439         self.next_control_block,440         bitmap,441         self.nvm3_key[0], self.nvm3_key[1], self.nvm3_key[2],442         self.max_link_data_len) = struct.unpack('<HHQbbbH', data)443        self.version = self.header & 0xf000 # Upper 4 bits of the header contain version number. Other bits are reserved for future use.444        self.next_control_block |= NVM3_CONTROL_BLOCK_KEY[args.protocol] & 0xF0000 # Only stored on two bytes. Adding prefix.445446        # Inspect bitmap.447        print('Bitmap: ', bin(bitmap))448449        if bitmap & (1 << CERTIFICATE_ON_DEVICE_BIT):450            self.certificate_on_device = True451            print('\tCertificate on device is required. NVM3 key: '452                  + hex(self.nvm3_key[CERTIFICATE_ON_DEVICE_BIT]))453        else:454            self.certificate_on_device = False455            print('\tCertificate on device is not required.')456        457        if bitmap & (1 << DEVICE_EC_KEY_BIT):458            self.device_ec_key_present = True459            print('\tEC key pair is present.')460        else:461            self.device_ec_key_present = False462            print('\tEC key pair is missing.')463464        if bitmap & (1 << STATIC_AUTH_DATA_BIT):465            self.static_auth_present = True466            print('\tStatic Authentication data is present.')467        else:468            self.static_auth_present = False469            print('\tStatic Authentication data is not present.')470        471        print('Maximum link data length: ', self.max_link_data_len)472473    def validate(self):474        # Expecting zero as the version number.475        if self.version != 0:476            raise Exception('Invalid Control Block! Version: ', self.version, '. Zero is expected.')477        else:478            print('Header: ', hex(self.header), '. Version: ', self.version)479        480        # Expecting only one control block. Therefore it should point onto itself.481        if self.next_control_block != NVM3_CONTROL_BLOCK_KEY[args.protocol]:482            raise Exception('Invalid Control Block! Next Control Block NVM3 key: ',483                            hex(self.next_control_block), '. ',484                            hex(NVM3_CONTROL_BLOCK_KEY[args.protocol]), ' is expected.')485        else:486            print('Next Control Block: ', hex(self.next_control_block), '\n\n')487        488        # Check the security of NVM3 ITS data.489        if self.device_ec_key_present and self.nvm3_key[DEVICE_EC_KEY_BIT] != 0:490            raise Exception('EC key pair is exposed. Any data stored in ITS should not be accessible!')491492        if self.static_auth_present and self.nvm3_key[STATIC_AUTH_DATA_BIT] != 0:493            raise Exception('Static Authentication data is exposed. Any data stored in ITS should not be accessible!')494495def subprocess_call(command):496    '''Handle subprocess calls.497498    Keyword arguments:499    command -- The CLI command to be called.500    Return values:501    process -- The CompletedProcess instance of the subprocess.502    '''503    try:504        process = subprocess.run(command,505                                 check=True,506                                 stdout=subprocess.PIPE,507                                 stderr=subprocess.STDOUT) # Merge stderr into stdout.508    except subprocess.CalledProcessError as e:509        raise Exception('[E: ', e.returncode, '] ', e.stdout)510    except:511        raise Exception('[E: ', e.returncode, '] ', e.stdout)512    return process513514def device_argument():515    '''Returns the device argument of a Simplicity Commander command.516    '''517    if args.serial is not None and args.ip is not None:518        raise Exception('At most one of [J-Link serial number] or [IP address] shall be defined.')519    elif args.serial is not None:520        return '--serialno ' + args.serial521    elif args.ip is not None:522        return '--ip ' + args.ip523    else:524        return ''525526class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter):527    pass528529def load_args():530    '''Parse command line arguments.'''531    parser = argparse.ArgumentParser(description=__doc__,532                                     formatter_class=CustomFormatter)533534    parser.epilog = (535        'examples:\n'536        '  %(prog)s                     Try to autodetect device.\n'537        '  %(prog)s --serial 440192051  Connect to device with the given J-Link serial.\n'538        '  %(prog)s --ip 192.168.0.143  Connect to device with the given IP address.')539540    parser.add_argument('-l', '--level',541                        default='0',542                        type=int,543                        help='The level authority that shall sign the request.')544    parser.add_argument('-v', '--validity',545                        default='365',546                        type=int,547                        help='The valid period of the certificate in days starting '\548                             'from the moment of signing.')549    parser.add_argument('-s', '--serial',550                        type=str,551                        help='J-Link serial number. Should not be given together with the IP address.')552    parser.add_argument('-i', '--ip',553                        type=str,554                        help='IP Address. Should not be given together with the J-Link serial number.')555    parser.add_argument('-p', '--protocol',556                        default='btmesh',557                        type=str.lower,558                        choices=['ble', 'btmesh'],559                        help='Determines which NVM3 region to use.')560561    args = parser.parse_args()562    return args563564if __name__ == '__main__':565    args = load_args()
...

server_v2.py

Source:server_v2.py

...380                    "target_name": project.config.target_name,381                    "profile_name": project.config.profile_name,382                    "logs": project.config.log_path,383                    "runner_parse_iteration": project._version,384                    "adapter_ready": project.adapter_probe(),385                }386                if project is not None387                else {}388            ),389            "timestamp": datetime.datetime.utcnow(),390            "error": None,391        },392        "dbt-osmosis-server": __name__,393    }394async def _adapter_heartbeat(runner: DbtProject):395    """Equivalent of a keepalive for adapters such as Snowflake"""396    await asyncio.sleep(60 * 30)397    while runner.adapter_probe():398        await asyncio.sleep(60 * 30)399def run_server(host="localhost", port=8581):400    uvicorn.run(401        "dbt_osmosis.core.server_v2:app",402        host=host,403        port=port,404        log_level="info",405        reload=False,406        workers=1,407    )408def test_server():409    """Some quick and dirty functional tests for the server"""410    import random411    import time...

Automation Testing Tutorials

Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.