Skip to content

Update configure_forensics.py #126

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
greg-aligiannis wants to merge 1 commit into from
Closed

Update configure_forensics.py #126

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
164 changes: 108 additions & 56 deletions forensic-archive/configure_forensics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,19 +1,9 @@
import boto3
import json
import os
import re
import argparse
from botocore.exceptions import ClientError
# -------------------------------
# Usage:
#
# python configure_forensics.py --bucket <bucket-name>
# -------------------------------


# -------------------------------
# Configuration
# -------------------------------
ROLE_PREFIXES = ["elastio-account-level-stack-", "elastio-account-level-sta-"]
ROLE_SUFFIXES = [
"awsS3ScanBgJob", "awsFsxOntapScanBgJob", "awsEfsScanBgJob", "awsEc2ScanBgJob",
Expand All @@ -23,12 +13,42 @@
]
POLICY_FILENAME = "policy.json"

# -------------------------------
# Helpers
# -------------------------------
def get_account_id():
return boto3.client("sts").get_caller_identity()["Account"]

def bucket_exists(bucket_name):
s3 = boto3.client('s3')
try:
s3.head_bucket(Bucket=bucket_name)
return True
except ClientError as e:
code = e.response['Error']['Code']
if code in ['404', 'NoSuchBucket']:
return False
elif code == '403':
print(f"⚠ Access denied to bucket: {bucket_name} (may exist in another account)")
return True
else:
raise

def create_bucket(bucket_name):
session = boto3.session.Session()
region = session.region_name or 'us-east-1'
s3 = session.client('s3')

try:
if region == 'us-east-1':
s3.create_bucket(Bucket=bucket_name)
else:
s3.create_bucket(
Bucket=bucket_name,
CreateBucketConfiguration={'LocationConstraint': region}
)
print(f"✅ Created bucket: {bucket_name} in region {region}")
except ClientError as e:
print(f"❌ Failed to create bucket: {bucket_name}\nReason: {e}")
raise

def get_verified_roles(account_id):
iam = boto3.client('iam')
paginator = iam.get_paginator('list_roles')
Expand Down Expand Up @@ -77,66 +97,92 @@ def save_policy_to_file(policy):
json.dump(policy, f, indent=4)
print(f"\n📁 Policy written to: {os.path.abspath(POLICY_FILENAME)}")

def update_batch_jobs(bucket_name):
env_vars_to_add = [
{'name': 'ELASTIO_FORENSIC_ANALYSIS_ARCHIVAL_ENABLED', 'value': 'TRUE'},
{'name': 'ELASTIO_FORENSIC_ANALYSIS_ARCHIVAL_S3_BUCKET_NAME', 'value': bucket_name}
]
def apply_policy_to_bucket(bucket_name, policy):
s3 = boto3.client('s3')
try:
s3.put_bucket_policy(Bucket=bucket_name, Policy=json.dumps(policy))
print(f"✅ Successfully applied policy to bucket: {bucket_name}")
except ClientError as e:
print(f"❌ Failed to apply policy to bucket: {bucket_name}\nReason: {e}")

def collect_elastio_job_names():
ec2 = boto3.client('ec2')
regions = [r['RegionName'] for r in ec2.describe_regions()['Regions']]
print(f"\n🌐 Found {len(regions)} AWS regions.")
job_names = set()

print("\n🔍 Collecting Elastio job definitions across all regions...")
for region in regions:
print(f"\n🌍 Checking region: {region}")
batch = boto3.client('batch', region_name=region)

try:
batch = boto3.client('batch', region_name=region)
response = batch.describe_job_definitions(status='ACTIVE')
job_definitions = response['jobDefinitions']
for jd in response.get('jobDefinitions', []):
name = jd['jobDefinitionName']
if name.startswith("elastio"):
job_names.add(name)
except ClientError as e:
print(f"⚠ Skipping region {region} due to error: {e}")
continue
print(f"⚠ Error collecting from region {region}: {e}")
return list(job_names), regions

if not job_definitions:
print("No active job definitions found.")
continue
def update_batch_jobs(bucket_name):
job_names, regions = collect_elastio_job_names()
if not job_names:
print("❌ No Elastio job definitions found.")
return

for job_def in job_definitions:
job_name = job_def['jobDefinitionName']
revision = job_def['revision']
env_vars_to_add = [
{'name': 'ELASTIO_FORENSIC_ANALYSIS_ARCHIVAL_ENABLED', 'value': 'TRUE'},
{'name': 'ELASTIO_FORENSIC_ANALYSIS_ARCHIVAL_S3_BUCKET_NAME', 'value': bucket_name}
]

print(f"\n🛠 Updating job definitions for {len(job_names)} jobs across {len(regions)} regions...")
for region in regions:
print(f"\n🌍 Region: {region}")
batch = boto3.client('batch', region_name=region)
for job_name in job_names:
try:
defs = batch.describe_job_definitions(jobDefinitionName=job_name, status='ACTIVE')
if not defs['jobDefinitions']:
continue

if job_name.startswith("elastio"):
print(f"🔧 Updating job: {job_name}:{revision}")
# Get the latest revision
latest_def = sorted(defs['jobDefinitions'], key=lambda d: d['revision'], reverse=True)[0]

container_props = job_def['containerProperties']
existing_env = container_props.get('environment', [])
# Use jobDefinitionName and find the matching revision
all_defs = batch.describe_job_definitions(jobDefinitionName=latest_def['jobDefinitionName'])['jobDefinitions']
full_def = next((d for d in all_defs if d['revision'] == latest_def['revision']), None)

if not full_def:
print(f"❌ Could not find full definition for {job_name}:{latest_def['revision']}")
continue

container_props = full_def['containerProperties']
existing_env = container_props.get('environment', [])
env_dict = {env['name']: env['value'] for env in existing_env}
for new_env in env_vars_to_add:
env_dict[new_env['name']] = new_env['value']
updated_env = [{'name': k, 'value': v} for k, v in env_dict.items()]
container_props['environment'] = updated_env

# Remove fields not accepted in registration
for field in ['jobDefinitionArn', 'revision', 'status']:
full_def.pop(field, None)

response = batch.register_job_definition(
jobDefinitionName=full_def['jobDefinitionName'],
type=full_def['type'],
parameters=full_def.get('parameters', {}),
containerProperties=container_props,
retryStrategy=full_def.get('retryStrategy', {}),
timeout=full_def.get('timeout', {}),
platformCapabilities=full_def.get('platformCapabilities', []),
tags=full_def.get('tags', {})
)
print(f"✅ Updated {job_name} -> {response['jobDefinitionArn']}")
except ClientError as e:
print(f"❌ Failed to update {job_name} in {region}: {e}")

print("\n🎉 Completed updating job definitions in all regions.")

sanitized_props = {
k: v for k, v in container_props.items() if k != 'networkConfiguration'
}
sanitized_props['environment'] = updated_env

try:
response = batch.register_job_definition(
jobDefinitionName=job_name,
type=job_def['type'],
containerProperties=sanitized_props
)
print(f"✅ Registered new revision: {response['jobDefinitionArn']}")
except ClientError as e:
print(f"❌ Failed to register {job_name} in {region}: {e}")

print("\n🎉 Multi-region job update completed.")

# -------------------------------
# Entry Point
# -------------------------------
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Configure S3 bucket policy and enable forensic archival.")
parser.add_argument("--bucket", type=str, default="elastio-forensic", help="S3 bucket name")
Expand All @@ -145,6 +191,12 @@ def update_batch_jobs(bucket_name):

print(f"\n🚀 Starting configuration for bucket: {bucket_name}")

if bucket_exists(bucket_name):
print(f"🪣 Bucket already exists: {bucket_name}")
else:
print(f"📦 Bucket {bucket_name} not found. Creating it...")
create_bucket(bucket_name)

account_id = get_account_id()
role_arns = get_verified_roles(account_id)

Expand All @@ -155,5 +207,5 @@ def update_batch_jobs(bucket_name):
print("\n✅ Generated policy:\n")
print(json.dumps(policy, indent=4))
save_policy_to_file(policy)

apply_policy_to_bucket(bucket_name, policy)
update_batch_jobs(bucket_name)