Merge pull request #16 from hellmarbecker/main

Add SASL authentication as an option

Author: ftisiot

README.md

@@ -24,7 +24,7 @@ pip install -r requirements.txt
 ## Usage
 
 The Python code can be run in bash with the following,
-in SSL security protocol:
+in ``SSL`` security protocol:
 ```bash
 python main.py \
  --security-protocol ssl \
@@ -35,7 +35,21 @@ python main.py \
  --nr-messages 0 \
  --max-waiting-time 0
 ```
-in PLAINTEXT security protocol:
+in ``SASL_SSL`` security protocol:
+```bash
+python main.py \
+ --security-protocol SASL_SSL \
+ --sasl-mechanism SCRAM-SHA-256 \
+ --username <USERNAME> \
+ --password <PASSWORD> \
+ --cert-folder ~/kafkaCerts/ \
+ --host kafka-<name>.aivencloud.com \
+ --port 13041 \
+ --topic-name pizza-orders \
+ --nr-messages 0 \
+ --max-waiting-time 0
+```
+in ``PLAINTEXT`` security protocol:
 ```bash
 python main.py \
  --security-protocol plaintext \
@@ -46,7 +60,7 @@ python main.py \
  --max-waiting-time 0
 ```
 Where
-* `security-protocol`: Security protocol for Kafka. PLAINTEXT or SSL are supported.
+* `security-protocol`: Security protocol for Kafka. ``PLAINTEXT``, ``SSL`` or ``SASL_SSL`` are supported.
 * `cert-folder`: points to the folder containing the Apache Kafka CA certificate, Access certificate and Access key (see [blog post](https://aiven.io/blog/create-your-own-data-stream-for-kafka-with-python-and-faker?utm_source=github&utm_medium=organic&utm_campaign=blog_art&utm_content=post) for more)
 * `host`: the Apache Kafka host
 * `port`: the Apache Kafka port

main.py

@@ -17,15 +17,17 @@
 MAX_ADDITIONAL_TOPPINGS_IN_PIZZA = 5
 
 
-
 # Creating a Faker instance and seeding to have the same results every time we execute the script
 fake = Faker()
 Faker.seed(4321)
 
 
 # function produce_msgs starts producing messages with Faker
 def produce_msgs(security_protocol='SSL',
+ sasl_mechanism='SCRAM-SHA-256',
  cert_folder = '~/kafka-pizza/',
+ username = '',
+ password = '',
  hostname = 'hostname',
  port = '1234',
  topic_name = 'pizza-orders',
@@ -49,6 +51,17 @@ def produce_msgs(security_protocol='SSL',
  value_serializer=lambda v: json.dumps(v).encode('ascii'),
  key_serializer=lambda v: json.dumps(v).encode('ascii')
  )
+ elif security_protocol.upper() == 'SASL_SSL':
+ producer = KafkaProducer(
+ bootstrap_servers=hostname + ':' + port,
+ security_protocol='SASL_SSL',
+ sasl_mechanism=sasl_mechanism,
+ ssl_cafile=cert_folder+'/ca.pem',
+ sasl_plain_username=username,
+ sasl_plain_password=password,
+ value_serializer=lambda v: json.dumps(v).encode('ascii'),
+ key_serializer=lambda v: json.dumps(v).encode('ascii')
+ )
  else:
  sys.exit("This security protocol is not supported!")
 
@@ -96,8 +109,11 @@ def produce_msgs(security_protocol='SSL',
 
 def main():
  parser = argparse.ArgumentParser()
- parser.add_argument('--security-protocol', help='Security protocol for Kafka (PLAINTEXT, SSL)', required=True)
- parser.add_argument('--cert-folder', help='Path to folder containing required Kafka certificates. Required --security-protocol equal SSL', required=False)
+ parser.add_argument('--security-protocol', help='Security protocol for Kafka (PLAINTEXT, SSL, SASL_SSL)', required=True)
+ parser.add_argument('--sasl-mechanism', help='SASL mechanism for Kafka (PLAIN, GSSAPI, OAUTHBEARER, SCRAM-SHA-256, SCRAM-SHA-512)', required=False)
+ parser.add_argument('--cert-folder', help='Path to folder containing required Kafka certificates. Required --security-protocol equal SSL or SASL_SSL', required=False)
+ parser.add_argument('--username', help='Username. Required if security-protocol is SASL_SSL', required=False)
+ parser.add_argument('--password', help='Password. Required if security-protocol is SASL_SSL', required=False)
  parser.add_argument('--host', help='Kafka Host (obtained from Aiven console)', required=True)
  parser.add_argument('--port', help='Kafka Port (obtained from Aiven console)', required=True)
  parser.add_argument('--topic-name', help='Topic Name', required=True)
@@ -106,19 +122,25 @@ def main():
  parser.add_argument('--subject', help='What type of content to produce (possible choices are [pizza, userbehaviour, stock, realstock, metric] pizza is the default', required=False)
  args = parser.parse_args()
  p_security_protocol = args.security_protocol
- p_cert_folder =args.cert_folder
- p_hostname =args.host
- p_port =args.port
- p_topic_name=args.topic_name
- p_subject=args.subject
+ p_cert_folder = args.cert_folder
+ p_username = args.username
+ p_password = args.password
+ p_sasl_mechanism = args.sasl_mechanism
+ p_hostname = args.host
+ p_port = args.port
+ p_topic_name= args.topic_name
+ p_subject = args.subject
  produce_msgs(security_protocol=p_security_protocol,
  cert_folder=p_cert_folder,
+ username=p_username,
+ password=p_password,
  hostname=p_hostname,
  port=p_port,
  topic_name=p_topic_name,
  nr_messages=int(args.nr_messages),
  max_waiting_time_in_sec=float(args.max_waiting_time),
- subject=p_subject
+ subject=p_subject,
+ sasl_mechanism=p_sasl_mechanism,
  )
  print(args.nr_messages)