<commit_before>#  Copyright 2015 Observable Networks
#
#  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.
import sys

from argparse import ArgumentParser
from datetime import datetime
from itertools import chain
from uuid import uuid4

import boto3

from .aggregation import aggregated_records
from .flowlogs_reader import FlowLogsReader, S3FlowLogsReader, SKIPDATA, NODATA

actions = {}


def action_print(reader, *args):
    """Simply print the Flow Log records to output."""
    arg_count = len(args)
    if arg_count == 0:
        stop_after = 0
    elif arg_count == 1:
        stop_after = int(args[0])
    else:
        raise RuntimeError("0 or 1 arguments expected for action 'print'")

    for i, record in enumerate(reader, 1):
        print(record.to_message())
        if i == stop_after:
            break


actions['print'] = action_print


def action_ipset(reader, *args):
    """Show the set of IPs seen in Flow Log records."""
    ip_set = set()
    for record in reader:
        if record.log_status in (SKIPDATA, NODATA):
            continue
        ip_set.add(record.srcaddr)
        ip_set.add(record.dstaddr)

    for ip in ip_set:
        print(ip)


actions['ipset'] = action_ipset


def action_findip(reader, *args):
    """Find Flow Log records involving a specific IP or IPs."""
    target_ips = set(args)
    for record in reader:
        if (record.srcaddr in target_ips) or (record.dstaddr in target_ips):
            print(record.to_message())


actions['findip'] = action_findip


def action_aggregate(reader, *args):
    """Aggregate flow records by 5-tuple and print a tab-separated stream"""
    all_aggregated = aggregated_records(reader)
    first_row = next(all_aggregated)
    keys = sorted(first_row.keys())
    print(*keys, sep='\t')

    # Join the first row with the rest of the rows and print them
    iterable = chain([first_row], all_aggregated)
    for item in iterable:
        print(*[item[k] for k in keys], sep='\t')


actions['aggregate'] = action_aggregate


def get_reader(args):
    kwargs = {}
    time_format = args.time_format

    if args.location_type == 'cwl':
        cls = FlowLogsReader
        client_type = 'logs'
    elif args.location_type == 's3':
        cls = S3FlowLogsReader
        client_type = 's3'

    if args.region:
        kwargs['region_name'] = args.region

    if args.profile:
        kwargs['profile_name'] = args.profile

    if args.start_time:
        kwargs['start_time'] = datetime.strptime(args.start_time, time_format)

    if args.end_time:
        kwargs['end_time'] = datetime.strptime(args.end_time, time_format)

    if args.location_type == 'cwl' and args.filter_pattern:
        kwargs['filter_pattern'] = args.filter_pattern

    if args.location_type == 's3' and args.include_accounts:
        kwargs['include_accounts'] = [
            x.strip() for x in args.include_accounts.split(',')
        ]

    if args.location_type == 's3' and args.include_regions:
        kwargs['include_regions'] = [
            x.strip() for x in args.include_regions.split(',')
        ]

    if args.thread_count:
        kwargs['thread_count'] = args.thread_count

    # Switch roles for access to another account
    if args.role_arn:
        assume_role_kwargs = {}
        assume_role_kwargs['RoleArn'] = args.role_arn<commit_msg>The parameter is only for S3<commit_after>#  Copyright 2015 Observable Networks
#
#  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.
import sys

from argparse import ArgumentParser
from datetime import datetime
from itertools import chain
from uuid import uuid4

import boto3

from .aggregation import aggregated_records
from .flowlogs_reader import FlowLogsReader, S3FlowLogsReader, SKIPDATA, NODATA

actions = {}


def action_print(reader, *args):
    """Simply print the Flow Log records to output."""
    arg_count = len(args)
    if arg_count == 0:
        stop_after = 0
    elif arg_count == 1:
        stop_after = int(args[0])
    else:
        raise RuntimeError("0 or 1 arguments expected for action 'print'")

    for i, record in enumerate(reader, 1):
        print(record.to_message())
        if i == stop_after:
            break


actions['print'] = action_print


def action_ipset(reader, *args):
    """Show the set of IPs seen in Flow Log records."""
    ip_set = set()
    for record in reader:
        if record.log_status in (SKIPDATA, NODATA):
            continue
        ip_set.add(record.srcaddr)
        ip_set.add(record.dstaddr)

    for ip in ip_set:
        print(ip)


actions['ipset'] = action_ipset


def action_findip(reader, *args):
    """Find Flow Log records involving a specific IP or IPs."""
    target_ips = set(args)
    for record in reader:
        if (record.srcaddr in target_ips) or (record.dstaddr in target_ips):
            print(record.to_message())


actions['findip'] = action_findip


def action_aggregate(reader, *args):
    """Aggregate flow records by 5-tuple and print a tab-separated stream"""
    all_aggregated = aggregated_records(reader)
    first_row = next(all_aggregated)
    keys = sorted(first_row.keys())
    print(*keys, sep='\t')

    # Join the first row with the rest of the rows and print them
    iterable = chain([first_row], all_aggregated)
    for item in iterable:
        print(*[item[k] for k in keys], sep='\t')


actions['aggregate'] = action_aggregate


def get_reader(args):
    kwargs = {}
    time_format = args.time_format

    if args.location_type == 'cwl':
        cls = FlowLogsReader
        client_type = 'logs'
    elif args.location_type == 's3':
        cls = S3FlowLogsReader
        client_type = 's3'

    if args.region:
        kwargs['region_name'] = args.region

    if args.profile:
        kwargs['profile_name'] = args.profile

    if args.start_time:
        kwargs['start_time'] = datetime.strptime(args.start_time, time_format)

    if args.end_time:
        kwargs['end_time'] = datetime.strptime(args.end_time, time_format)

    if args.location_type == 'cwl' and args.filter_pattern:
        kwargs['filter_pattern'] = args.filter_pattern

    if args.location_type == 's3' and args.include_accounts:
        kwargs['include_accounts'] = [
            x.strip() for x in args.include_accounts.split(',')
        ]

    if args.location_type == 's3' and args.include_regions:
        kwargs['include_regions'] = [
            x.strip() for x in args.include_regions.split(',')
        ]

    if args.location_type == 's3' and args.thread_count:
        kwargs['thread_count'] = args.thread_count

    # Switch roles for access to another account
    if args.role_arn:
        assume_role_kwargs = {}
        assume_role_kwargs['RoleArn'] = args.role_arn
{'fraction_non_alphanumeric': 0.098828125, 'fraction_numerical': 0.006510416666666667, 'mean_word_length': 3.087812666311868, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6482413', 'n_tokens_mistral': 2364, 'n_tokens_neox': 2277, 'n_words': 719}