Recent, mi-am dorit foarte mult să găsesc o modalitate de a construi un API care să preia o adresă URL și să salveze o captură de ecran.

Cazul meu de utilizare inițial a fost simplu: dacă analizam e-mailurile de phishing, doream o modalitate ușoară de a obține o captură de ecran a adresei URL către care e-mailul încerca să-și direcționeze țintele.

Pentru a construi acest lucru, am folosit Terraform pentru a crea toată infrastructura necesară pentru a o configura în AWS, folosind Selenium, chromedriver și Chrome fără cap pentru a obține capturile de ecran.

Diagramă de nivel înalt care ilustrează ce va fi construit în AWS
Diagramă de nivel înalt care ilustrează ce va fi construit în AWS de Terraform

Notă: toate eșantioanele de cod provin de la PowerShell, așa că vă rugăm să scuzați notația „. ”.

Cerințe

  • Un cont AWS
  • Binar Terraform
  • Bucket S3 existent pentru a stoca starea Terraform (https://www.terraform.io/docs/backends/types/s3.html)
  • Utilizator AWS IAM și cheie de acces create cu permisiuni adecvate (acces programatic, grup administrativ) pentru utilizarea Terraform

Cum se configurează proiectul

Creați noul dvs. director și inițializați Terraform astfel:

mkdir .screenshot-service
cd .screenshot-service
.terraform init

Configurați furnizorul AWS

Creați un fișier numit provider.tf în rădăcina directorului de proiect. Apoi configurați cu valorile corespunzătoare pentru cheia de acces AWS și cheia secretă, precum și numele unui bucket S3 existent care va fi utilizat pentru a stoca fișierul de stare Terraform.

provider "aws" {
  region = "us-east-1"
  
  access_key = "ACCESSKEY"
  secret_key = "SECRETKEY"
}

terraform {
  backend "s3" {
    bucket = "EXISTING_BUCKET"
    region = "us-east-1"
    key = "KEYFORSTATE"
    access_key = "ACCESSKEY"
    secret_key = "SECRETKEY"
    encrypt = "true"
  }
}

Configurați cupa S3

Vom folosi o cupă S3 pentru a stoca toate capturile de ecran. Pentru a configura serviciul S3, creați un nou fișier numit în rădăcina proiectului dvs. s3.tf și adăugați următoarele:

resource "aws_s3_bucket" "screenshot_bucket" {
  bucket        = "STORAGE_BUCKET_NAME"
  force_destroy = true
  acl = "public-read"

  versioning {
    enabled = false
  }
}

Creați stratul Lambda

Să începem prin crearea stratului lambda care va conține binarele necesare. Mai întâi, din rădăcina proiectului, creați un folder numit chromedriver_layer: mkdir .chromedriver_layer.

Apoi, descărcați binarul crom și râul:

cd .chromedriver_layer
wget https://chromedriver.storage.googleapis.com/2.41/chromedriver_linux64.zip -OutFile .chromedriver.zip
wget https://github.com/adieuadieu/serverless-chrome/releases/download/v1.0.0-54/stable-headless-chromium-amazonlinux-2017-03.zip -OutFile .headless-chromium.zip
Expand-Archive .headless-chromium.zip
rm *.zip

În cele din urmă, trebuie să aranjăm acest lucru frumos pentru Terraform:

cd ..
Compress-Archive .chromedriver_layer -DestinationPath chromedriver_layer.zip

Cum se configurează Lambda

Infrastructura Lambda

Creați un fișier numit lambda.tf în rădăcina directorului de proiect. În primul rând, vom crea rolul de execuție necesar funcției noastre:

resource "aws_iam_role" "lambda_exec_role" {
  name        = "lambda_exec_role"
  description = "Execution role for Lambda functions"

  assume_role_policy = <<EOF
{
        "Version"  : "2012-10-17",
        "Statement": [
            {
                "Action"   : "sts:AssumeRole",
                "Principal": {  
                    "Service": "lambda.amazonaws.com"
                },
                "Effect": "Allow",
                "Sid"   : ""
            }
        ]
}
EOF
}

Apoi, vom adăuga câteva politici la rolul de execuție pe care l-am creat, care va permite funcției noastre să acceseze serviciile necesare:

resource "aws_iam_role_policy" "lambda_logging" {
  name = "lambda_logging"

  role = aws_iam_role.lambda_exec_role.id

  policy = <<EOF
{
    "Version"  : "2012-10-17",
    "Statement": [
        {
            "Effect"  : "Allow",
            "Resource": "*",
            "Action"  : [
                "logs:CreateLogStream",
                "logs:PutLogEvents",
                "logs:CreateLogGroup"
            ]
        }
    ]
}
EOF
}

resource "aws_iam_role_policy" "lambda_s3_access" {
  name = "lambda_s3_access"

  role = aws_iam_role.lambda_exec_role.id

  # TODO: Change resource to be more restrictive
  policy = <<EOF
{
  "Version"  : "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": [
        "s3:ListBuckets",
        "s3:PutObject",
        "s3:PutObjectAcl",
        "s3:GetObjectAcl"
      ],
      "Resource": ["*"]
    }
  ]
}
EOF
}

Acolo, acum funcția noastră va putea accesa S3 și conecta la CloudWatch. Să ne definim funcția:

resource "aws_lambda_function" "take_screenshot" {
  filename      = "./screenshot-service.zip"
  function_name = "take_screenshot"
  role          = aws_iam_role.lambda_exec_role.arn
  handler       = "screenshot-service.handler"
  runtime       = "python3.7"

  source_code_hash = filebase64sha256("./screenshot-service.zip")
  timeout          = 600
  memory_size      = 512 
  layers = ["${aws_lambda_layer_version.chromedriver_layer.arn}"]

  environment {
    variables = {
      s3_bucket = "${aws_s3_bucket.screenshot_bucket.bucket}"
    }
  }
}

Codul de mai sus specifică faptul că încărcăm un pachet de funcții lambda folosind un runtime Python 3.7 și că funcția care va fi numită se numește “handler”.

Am setat timpul de expirare la 600 de secunde, dar nu ezitați să schimbați acest lucru după cum doriți. De asemenea, nu ezitați să jucați cu memory_size – pentru mine, acest lucru a dus la capturi de ecran super rapide.

De asemenea, am setat o variabilă de mediu numită s3_bucket care va fi transmis funcției, care conține numele cupei utilizate pentru a stoca captura de ecran.

Funcția Lambda în sine

Creați un folder numit lambda în rădăcina directorului proiectului și creați un fișier numit screenshot-service.py în acel folder.

Adăugați în fișier următoarele configurații de import și jurnal:

#!/usr/bin/env python
# -*- coding utf-8 -*-

import json
import logging
from urllib.parse import urlparse, unquote # TODO: Can I use urllib3?
from selenium import webdriver
from datetime import datetime
import os
from shutil import copyfile
import boto3
import stat
import urllib.request
import tldextract

# Configure logging
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)

Apoi, vom crea o funcție care va copia binarele din stratul nostru lambda și le va face executabile:

def configure_binaries():
    """Copy the binary files from the lambda layer to /tmp and make them executable"""
    copyfile("/opt/chromedriver", "/tmp/chromedriver")
    copyfile("/opt/headless-chromium", "/tmp/headless-chromium")

    os.chmod("/tmp/chromedriver", 755)
    os.chmod("/tmp/headless-chromium", 755)

Apoi, vom crea funcția care va face captura de ecran a domeniului furnizat. Vom transmite adresa URL și numele cupei S3.

Vom adăuga un parametru opțional care să permită setarea titlului imaginii de către utilizator. Captura de ecran este făcută de Selenium automatizând browserul Chrome fără cap pe care l-am descărcat.

def get_screenshot(url, s3_bucket, screenshot_title = None):     
    configure_binaries()

    chrome_options = webdriver.ChromeOptions()
    chrome_options.add_argument('--headless')
    chrome_options.add_argument("disable-infobars")
    chrome_options.add_argument("enable-automation")
    
    chrome_options.add_argument('--no-sandbox')
    chrome_options.add_argument('--disable-gpu')
    chrome_options.add_argument('--window-size=1280x1696')
    chrome_options.add_argument('--user-data-dir=/tmp/user-data')
    chrome_options.add_argument('--hide-scrollbars')
    chrome_options.add_argument('--enable-logging')
    chrome_options.add_argument('--log-level=0')
    chrome_options.add_argument('--disable-dev-shm-usage')
    chrome_options.add_argument('--v=99')
    chrome_options.add_argument('--single-process')
    chrome_options.add_argument('--data-path=/tmp/data-path')
    chrome_options.add_argument('--ignore-certificate-errors')
    chrome_options.add_argument('--homedir=/tmp')
    chrome_options.add_argument('--disk-cache-dir=/tmp/cache-dir')
    chrome_options.add_argument(
        'user-agent=Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/61.0.3163.100 Safari/537.36')
    chrome_options.binary_location = "/tmp/headless-chromium"

    if screenshot_title is None: 
        ext = tldextract.extract(url)
        domain = f"{''.join(ext[:2])}:{urlparse(url).port}.{ext[2]}"
        screenshot_title = f"{domain}_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}"
    logger.debug(f"Screenshot title: {screenshot_title}")

    with webdriver.Chrome(chrome_options=chrome_options, executable_path="/tmp/chromedriver", service_log_path="/tmp/selenium.log") as driver: 
        driver.set_window_size(1024, 768)
        
        logger.info(f"Obtaining screenshot for {url}")
        driver.get(url)     
        
        driver.save_screenshot(f"/tmp/{screenshot_title}.png") # TODO: Delete the screenshot after
        logger.info(f"Uploading /tmp/{screenshot_title}.png to S3 bucket {s3_bucket}/{screenshot_title}.png")
        s3 = boto3.client("s3")
        s3.upload_file(f"/tmp/{screenshot_title}.png", s3_bucket, f"{screenshot_title}.png", ExtraArgs={'ContentType': 'image/png', 'ACL': 'public-read'})
    return f"https://{s3_bucket}.s3.amazonaws.com/{screenshot_title}.png"

În sfârșit, să creăm handler-ul nostru, care va fi invocat atunci când API Gateway primește o solicitare legitimă:

def handler(event, context): 
    logger.debug("## ENVIRONMENT VARIABLES ##")
    logger.debug(os.environ)
    logger.debug("## EVENT ##")
    logger.debug(event)

    bucket_name = os.environ["s3_bucket"]
    logger.debug(f"bucket_name: {bucket_name}")

    logger.info("Validating url")  

    if event["httpMethod"] == "GET":
        if event["queryStringParameters"]:
            try:
                url = event["queryStringParameters"]["url"]
            except Exception as e:
                logger.error(e)
                raise e
        else:
            return {
                "statusCode": 400,
                "body": json.dumps("No URL provided...")
            }
    elif event["httpMethod"] == "POST":
        if event["body"]:
            try:
                body = json.loads(event["body"])
                url = body["url"]
            except Exception as e:
                logger.error(e)
                raise e
        else:
            return {
                "statusCode": 400,
                "body": json.dumps("No URL provided...")
            }
    else:
        return {
            "statusCode": 405,
            "body": json.dumps(f"Invalid HTTP Method {event['httpMethod']} supplied")
        }

    logger.info(f"Decoding {url}")
    url = unquote(url)

    logger.info(f"Parsing {url}")
    try: 
        parsed_url = urlparse(url)
        if parsed_url.scheme != "http" and parsed_url.scheme != "https":
            logger.info("No valid scheme found, defaulting to http://")
            parsed_url = urlparse(f"http://{url}")
        if parsed_url.port is None:
            if parsed_url.scheme == "http":
                parsed_url = urlparse(f"{parsed_url.geturl()}:80")
            elif parsed_url.scheme == "https":
                parsed_url = urlparse(f"{parsed_url.geturl()}:443")

    except Exception as e: 
        logger.error(e)
        raise e
    
    logger.info("Getting screenshot")
    try: 
        screenshot_url = get_screenshot(parsed_url.geturl(), bucket_name) # TODO: Variable!
    except Exception as e:  
        logger.error(e)
        raise e

    response_body = {
        "message": f"Successfully captured screenshot of {parsed_url.geturl()}",
        "screenshot_url": screenshot_url
    }

    return {
        "statusCode": 200,
        "body"      : json.dumps(response_body)
    }

Apoi, trebuie să instalăm toate pachetele pe care funcția lambda le folosește în lambda director, deoarece aceste pachete nu sunt instalate implicit în AWS.

Apoi, trebuie să creăm arhiva zip (odată creată, Terraform va continua să o actualizeze dacă modificați codul):

cd .lambda
pip install selenium tldextract -t .
cd ..
Compress-Archive .lambda -DestinationPath .screenshot-service.zip

Cum se configurează API Gateway

Creați un fișier numit apigw.tf în rădăcina directorului de proiect. Mai întâi, vom configura API-ul REST:

resource "aws_api_gateway_rest_api" "screenshot_api" {
  name        = "screenshot_api"
  description = "Lambda-powered screenshot API"
  depends_on = [
    aws_lambda_function.take_screenshot
  ]
}

Acest API va fi utilizat pentru a direcționa toate cererile care sunt făcute pentru serviciul de captură de ecran. Noi folosim depends_on caracteristică pentru a se asigura că gateway-ul și componentele sale conexe sunt create numai după se creează funcția lambda.

Apoi, să creăm resursa API Gateway pentru funcția lambda:

resource "aws_api_gateway_resource" "screenshot_api_gateway" {
  path_part   = "screenshot"
  parent_id   = aws_api_gateway_rest_api.screenshot_api.root_resource_id
  rest_api_id = aws_api_gateway_rest_api.screenshot_api.id
}

Am definit acum o resursă care va răspunde la /screenshot punct final pentru serviciul API.

În continuare vom crea o etapă pentru API. O etapă este un mod elegant de a numi implementarea API-ului nostru. Puteți configura stocarea în cache, înregistrarea în funcțiune, limitarea cererii și multe altele utilizând o etapă.

resource "aws_api_gateway_stage" "prod_stage" {
  stage_name = "prod"
  rest_api_id = aws_api_gateway_rest_api.screenshot_api.id
  deployment_id = aws_api_gateway_deployment.api_gateway_deployment_get.id
}

În continuare vom crea o cheie API și un plan de utilizare legat de etapa noastră, astfel încât numai utilizatorii cu cunoștințe despre cheie să poată utiliza acest serviciu. (Notă: Dacă doriți acest lucru accesibil publicului, săriți peste acest pas.)

resource "aws_api_gateway_usage_plan" "apigw_usage_plan" {
  name = "apigw_usage_plan"

  api_stages {
    api_id = aws_api_gateway_rest_api.screenshot_api.id
    stage = aws_api_gateway_stage.prod_stage.stage_name
  }
}

resource "aws_api_gateway_usage_plan_key" "apigw_usage_plan_key" {
  key_id = aws_api_gateway_api_key.apigw_prod_key.id
  key_type = "API_KEY"
  usage_plan_id = aws_api_gateway_usage_plan.apigw_usage_plan.id
}

resource "aws_api_gateway_api_key" "apigw_prod_key" {
  name = "prod_key"
}

Să configurăm acum API-ul pentru a răspunde fie la a OBȚINE sau POST solicitați dacă este furnizată o cheie Gateway API validă (setați valoarea la false dacă doriți ca metoda să fie deschisă publicului):

resource "aws_api_gateway_method" "take_screenshot_get" {
  rest_api_id   = aws_api_gateway_rest_api.screenshot_api.id
  resource_id   = aws_api_gateway_resource.screenshot_api_gateway.id
  http_method   = "GET"
  authorization = "NONE"
  api_key_required = true
}

resource "aws_api_gateway_method" "take_screenshot_post" {
  rest_api_id   = aws_api_gateway_rest_api.screenshot_api.id
  resource_id   = aws_api_gateway_resource.screenshot_api_gateway.id
  http_method   = "POST"
  authorization = "NONE"
  api_key_required = true
}

Acum trebuie să acordăm API Gateway permisiunea de a invoca funcția lambda pe care am creat-o:

resource "aws_lambda_permission" "apigw" {
  statement_id  = "AllowAPIGatewayInvoke"
  action        = "lambda:InvokeFunction"
  function_name = aws_lambda_function.take_screenshot.arn
  principal     = "apigateway.amazonaws.com"

  source_arn = "${aws_api_gateway_rest_api.screenshot_api.execution_arn}/*/*/*"
}

Excelent, acum avem permisiunile corespunzătoare. Să configurăm integrarea noastră cu funcția lambda:

resource "aws_api_gateway_integration" "lambda_integration_get" {
  depends_on = [
    aws_lambda_permission.apigw
  ]
  rest_api_id = aws_api_gateway_rest_api.screenshot_api.id
  resource_id = aws_api_gateway_method.take_screenshot_get.resource_id
  http_method = aws_api_gateway_method.take_screenshot_get.http_method

  integration_http_method = "POST" # https://github.com/hashicorp/terraform/issues/9271 Lambda requires POST as the integration type
  type                    = "AWS_PROXY"
  uri                     = aws_lambda_function.take_screenshot.invoke_arn
}

resource "aws_api_gateway_integration" "lambda_integration_post" {
  depends_on = [
    aws_lambda_permission.apigw
  ]
  rest_api_id = aws_api_gateway_rest_api.screenshot_api.id
  resource_id = aws_api_gateway_method.take_screenshot_post.resource_id
  http_method = aws_api_gateway_method.take_screenshot_post.http_method

  integration_http_method = "POST" # https://github.com/hashicorp/terraform/issues/9271 Lambda requires POST as the integration type
  type                    = "AWS_PROXY"
  uri                     = aws_lambda_function.take_screenshot.invoke_arn
}

Această integrare spune API Gateway ce funcție lambda trebuie să invoce atunci când primește o cerere la punctul final și metoda HTTP.

Aproape am terminat cu gateway-ul, promit. Ca ultim pas, să ne asigurăm că API-ul nostru poate trimite jurnale către CloudWatch:

resource "aws_api_gateway_account" "apigw_account" {
  cloudwatch_role_arn = aws_iam_role.apigw_cloudwatch.arn
}

resource "aws_iam_role" "apigw_cloudwatch" {
  # https://gist.github.com/edonosotti/6e826a70c2712d024b730f61d8b8edfc
  name = "api_gateway_cloudwatch_global"

  assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "",
      "Effect": "Allow",
      "Principal": {
        "Service": "apigateway.amazonaws.com"
      },
      "Action": "sts:AssumeRole"
    }
  ]
}
EOF
}

resource "aws_iam_role_policy" "apigw_cloudwatch" {
  name = "default"
  role = aws_iam_role.apigw_cloudwatch.id

  policy = <<EOF
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "logs:CreateLogGroup",
                "logs:CreateLogStream",
                "logs:DescribeLogGroups",
                "logs:DescribeLogStreams",
                "logs:PutLogEvents",
                "logs:GetLogEvents",
                "logs:FilterLogEvents"
            ],
            "Resource": "*"
        }
    ]
}
EOF
}

Acum am acordat API Gateway permisiunile necesare pentru a scrie jurnale în CloudWatch.

Nu în ultimul rând, implementăm API-ul nostru. Folosim depends_on pentru a vă asigura că implementarea are loc după crearea tuturor dependențelor.

resource "aws_api_gateway_deployment" "api_gateway_deployment_get" {
  depends_on = [aws_api_gateway_integration.lambda_integration_get,  aws_api_gateway_method.take_screenshot_get, aws_api_gateway_integration.lambda_integration_post, aws_api_gateway_method.take_screenshot_post]

  rest_api_id = aws_api_gateway_rest_api.screenshot_api.id
}

Ambalaj Lambda

În main.tf, adăugați următoarele:

data "archive_file" "screenshot_service_zip" {
  type        = "zip"
  source_dir  = "./lambda"
  output_path = "./screenshot-service.zip"
}

data "archive_file" "screenshot_service_layer_zip" {
  type = "zip"
  source_dir = "./chromedriver_layer"
  output_path = "./chromedriver_lambda_layer.zip"
}

Ieșiri

Creați un fișier numit output.tf în rădăcina directorului de proiect și adăugați următoarele:

output "api_gateway_url" {
  value = "${aws_api_gateway_stage.prod_stage.invoke_url}/${aws_api_gateway_resource.screenshot_api_gateway.path_part}"
}

output "api_key" {
  value = aws_api_gateway_api_key.apigw_prod_key.value
}

Acum, odată ce ai fugit .terraform apply veți obține rezultate cu adresa URL a API-ului și cheia API asociată.

Felicitări! Acum aveți un serviciu de capturi de ecran funcțional. Pentru a vizualiza codul pe care l-am folosit, nu ezitați să verificați Github repertoriu.